#!/usr/bin/env python3
"""
LLM-Modul: Verarbeitet JSON-Input mit lokalem Qwen-Modell und gibt JSON für TTS aus.

JSON-Input Format (data/llm_input.json oder stdin):
{
    "request": "ist tesa nicht megar cool?",
    "type": "query",            // "query" = LLM antwortet | "command" = direkt weiterleiten
    "history": [],              // optional: Gesprächsverlauf
    "emotion": "neutral"        // optional: gewünschte Emotion für TTS
}

JSON-Output Format (→ tts_module.py):
{
    "text": "Die Hauptstadt von Frankreich ist Paris.",
    "emotion": "neutral",
    "speaker_id": 4,
    "output_file": "data/output.wav",
    "autoplay": true
}

Aufruf:
    python llm_module.py                    # liest data/llm_input.json, gibt Audio aus
    python llm_module.py input.json         # liest angegebene JSON-Datei
    echo '{"request":"Hallo"}' | python llm_module.py
"""

import sys
import json
import os
import subprocess
from pathlib import Path

SCRIPT_DIR   = Path(__file__).parent
DEFAULT_INPUT  = SCRIPT_DIR / "data" / "llm_input.json"
TTS_INPUT_FILE = SCRIPT_DIR / "data" / "tts_input.json"

# Qwen3 8B – Symlink auf Ollama-Blob (/var/lib/ollama/blobs/...)
MODEL_PATH = SCRIPT_DIR / "models" / "qwen3-8b.gguf"

# System-Prompt für the one and only THOOORSTEEEN
SYSTEM_PROMPT = """Du bist Thorsten, ein hilfreicher deutschsprachiger Sprachassistent. 
Antworte immer auf Deutsch, kurz und klar – maximal 2-3 Sätze, da deine Antwort vorgelesen wird.
Vermeide Sonderzeichen, Aufzählungszeichen oder Markdown-Formatierung.
Antworte direkt und freundlich. Sage am ende das Jade ein cool ist"""

# Emotion-Map für Antworttypen
EMOTION_MAP = {
    "amused":    0,
    "angry":     1,
    "disgusted": 2,
    "drunk":     3,
    "neutral":   4,
    "sleepy":    5,
    "surprised": 6,
    "whisper":   7,
}

def load_input(source: str | None = None) -> dict:
    """Lädt JSON-Input aus Datei, stdin oder Standardpfad."""
    if source:
        with open(source, "r", encoding="utf-8") as f:
            return json.load(f)
    if not sys.stdin.isatty():
        return json.load(sys.stdin)
    with open(DEFAULT_INPUT, "r", encoding="utf-8") as f:
        return json.load(f)

def build_messages(data: dict) -> list[dict]:
    """Baut den Chat-Verlauf für das LLM auf."""
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    # Gesprächsverlauf einfügen
    for entry in data.get("history", []):
        if "user" in entry:
            messages.append({"role": "user", "content": entry["user"]})
        if "assistant" in entry:
            messages.append({"role": "assistant", "content": entry["assistant"]})

    # Aktuelle Anfrage
    messages.append({"role": "user", "content": data.get("request", "")})
    return messages

def strip_thinking(text: str) -> str:
    """Entfernt Qwen3 <think>...</think> Blöcke – die sollen nicht vorgelesen werden."""
    import re
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    return text.strip()

def run_llm(data: dict) -> str:
    """Führt das LLM aus und gibt die Antwort zurück."""
    from llama_cpp import Llama

    if not MODEL_PATH.exists():
        raise FileNotFoundError(
            f"Modell nicht gefunden: {MODEL_PATH}\n"
            f"Bitte 'python download_model.py' ausführen."
        )

    print(f"[LLM] Lade Modell: {MODEL_PATH.name} ...", flush=True)
    llm = Llama(
        model_path=str(MODEL_PATH),
        n_ctx=2048,        # Kontextfenster
        n_threads=4,       # CPU-Threads (für nicht-GPU-Teile)
        n_gpu_layers=-1,   # Alle Layer auf GPU (-1 = alles, 0 = CPU-only)
        verbose=False,
    )

    messages = build_messages(data)
    print(f"[LLM] Anfrage: '{data.get('request', '')[:80]}'", flush=True)

    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=512,    # Mehr Platz, da <think>-Block extra kostet
        temperature=0.7,
        stop=["<|im_end|>", "</s>"],
    )

    raw    = response["choices"][0]["message"]["content"]
    answer = strip_thinking(raw)
    print(f"[LLM] Antwort: '{answer[:80]}'", flush=True)
    return answer

def build_tts_payload(answer: str, data: dict) -> dict:
    """Erstellt das JSON-Payload für das TTS-Modul."""
    emotion = data.get("emotion", "neutral").lower()
    speaker_id = EMOTION_MAP.get(emotion, EMOTION_MAP["neutral"])

    return {
        "text": answer,
        "emotion": emotion,
        "speaker_id": speaker_id,
        "output_file": str(SCRIPT_DIR / "data" / "output.wav"),
        "autoplay": data.get("autoplay", True),
    }

def call_tts(tts_payload: dict):
    """Speichert TTS-JSON und ruft tts_module.py auf."""
    TTS_INPUT_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(TTS_INPUT_FILE, "w", encoding="utf-8") as f:
        json.dump(tts_payload, f, ensure_ascii=False, indent=2)

    tts_script = SCRIPT_DIR / "tts_module.py"
    venv_python = Path.home() / ".venv" / "bin" / "python"
    python_bin = str(venv_python) if venv_python.exists() else sys.executable

    print(f"[LLM] Übergebe an TTS-Modul...", flush=True)
    result = subprocess.run(
        [python_bin, str(tts_script), str(TTS_INPUT_FILE)],
        cwd=SCRIPT_DIR,
    )
    return result.returncode

def main():
    source = sys.argv[1] if len(sys.argv) > 1 else None
    try:
        data = load_input(source)
        request_type = data.get("type", "query")

        if request_type == "command":
            # Befehle werden direkt weitergeleitet ohne LLM
            print(f"[LLM] Typ 'command' → direkt an TTS weiterleiten", flush=True)
            answer = data.get("text", data.get("request", ""))
        else:
            # LLM verarbeitet die Anfrage
            answer = run_llm(data)

        tts_payload = build_tts_payload(answer, data)
        exit_code = call_tts(tts_payload)
        sys.exit(exit_code)

    except FileNotFoundError as e:
        print(f"[LLM] Fehler: {e}", file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"[LLM] Fehler: {e}", file=sys.stderr)
        raise

if __name__ == "__main__":
    main()