import gradio as gr
from llama_cpp import Llama

# --- 1. SETUP & DATEN ---

# Das lokale Modell laden (dein angegebener Pfad)
llm = Llama.from_pretrained(
    repo_id="simonper/Llama-3.2-1B-bnb-4bit_finetome-100k_gguf_4bit",
    filename="Llama-3.2-1B.Q4_K_M.gguf",
    n_ctx=4096,      # Context window
    n_threads=8,     # CPU threads (anpassbar)
    n_gpu_layers=0,  # 0 für CPU, höher setzen wenn du GPU hast
)

# Die Datenbank der Rätsel (Szenario = Öffentlich, Lösung = Versteckt)
RIDDLES = {
    "Der tote Mann": {
        "scenario": "Ein Mann liegt tot auf einer Wiese. Neben ihm liegt ein ungeöffnetes Paket. Es sind keine Fußspuren zu sehen. Wie ist er gestorben?",
        "solution": "Der Mann ist aus einem Flugzeug gesprungen. Das Paket war sein Fallschirm, der sich nicht geöffnet hat."
    },
    "Der Barman": {
        "scenario": "Ein Mann geht in eine Bar und bestellt ein Glas Wasser. Der Barkeeper zieht eine Waffe und richtet sie auf ihn. Der Mann sagt 'Danke' und geht. Warum?",
        "solution": "Der Mann hatte Schluckauf. Der Schreck durch die Waffe hat ihn geheilt, weshalb er dankbar war."
    },
    "Die Hütte": {
        "scenario": "Zwei Männer sind in einer Hütte im Wald. Einer ist tot. Die Hütte ist nicht abgebrannt, aber der tote Mann ist verkohlt. Was ist passiert?",
        "solution": "Die Hütte ist die Kabine eines abgestürzten Flugzeugs. Der Mann starb beim Absturzfeuer."
    },
    "Der Koffer": {
        "scenario": "Eine Frau öffnet ihren Koffer und findet einen toten Mann darin. Sie wird nicht verhaftet und hat keine Angst. Warum?",
        "solution": "Der 'Koffer' ist eigentlich ein Sarg. Die Frau ist auf einer Beerdigung."
    }
}

# --- 2. HELPER FUNCTIONS ---

def build_prompt(system_message: str, history: list[dict], user_message: str) -> str:
    """
    Konstruiert den Prompt für das Llama Modell.
    Da wir Llama-3 nutzen, ist ein Format wichtig, das System-Instruktionen klar trennt.
    """
    lines = []

    # System Prompt Injection
    if system_message:
        lines.append(f"System: {system_message}\n")

    # History aufbauen
    for turn in history:
        role = turn["role"]
        content = turn["content"]
        if role == "user":
            lines.append(f"User: {content}")
        elif role == "assistant":
            lines.append(f"Assistant: {content}")

    # Aktuelle Nachricht
    lines.append(f"User: {user_message}")
    lines.append("Assistant:")

    return "\n".join(lines)


def respond(
    message,
    history: list[dict[str, str]],
    system_message_dummy, # Wir ignorieren den Input vom Slider und bauen unseren eigenen
    max_tokens,
    temperature,
    top_p,
    selected_riddle, # Das kommt vom Dropdown
):
    # 1. Das aktuelle Rätsel laden
    current_game = RIDDLES[selected_riddle]
    
    # 2. Den "Game Master" System Prompt bauen (Hier passiert das Context Learning)
    # Wir injizieren die Lösung direkt in den Kontext des Modells
    game_master_prompt = (
        f"Du bist der Spielleiter eines Laterale-Denk-Rätsels (Black Stories). "
        f"AKTUELLES SZENARIO: '{current_game['scenario']}' "
        f"VERSTECKTE LÖSUNG (Der User kennt diese NICHT!): '{current_game['solution']}' "
        f"\nREGELN FÜR DICH:\n"
        f"1. Analysiere die Frage des Users logisch in Bezug auf die versteckte Lösung.\n"
        f"2. Antworte AUSSCHLIESSLICH mit: 'Ja', 'Nein', oder 'Irrelevant'.\n"
        f"3. Wenn der User die Lösung errät, sage: 'KORREKT! Du hast es gelöst: [Erklärung]'.\n"
        f"4. Gib KEINE Hinweise."
    )

    # 3. Prompt zusammenbauen
    prompt = build_prompt(game_master_prompt, history, message)

    # 4. Llama.cpp aufrufen
    output = llm(
        prompt,
        max_tokens=int(max_tokens),
        temperature=float(temperature),
        top_p=float(top_p),
        stop=["User:", "System:"],
        echo=False # Wichtig: Den Prompt nicht wiederholen
    )

    reply = output["choices"][0]["text"].strip()
    return reply


# --- 3. UI AUFBAU (GRADIO BLOCKS) ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🕵️ Der Ja/Nein Detektiv (Llama GGUF Edition)")
    gr.Markdown(
        "**Anleitung:** Ich denke an eine seltsame Situation. "
        "Stelle mir Ja/Nein-Fragen, um herauszufinden, was passiert ist!"
    )

    # Das Dropdown Menü für die Rätsel-Auswahl
    with gr.Row():
        riddle_select = gr.Dropdown(
            choices=list(RIDDLES.keys()), 
            value="Der tote Mann", 
            label="Wähle einen Fall", 
            interactive=True
        )
    
    # Textfeld, das das Szenario anzeigt
    scenario_display = gr.Textbox(
        label="Das Szenario (Dein Hinweis)", 
        value=RIDDLES["Der tote Mann"]["scenario"],
        interactive=False 
    )

    # Logik: Wenn Dropdown geändert wird, Textbox aktualisieren
    def update_scenario(choice):
        return RIDDLES[choice]["scenario"]
    
    riddle_select.change(fn=update_scenario, inputs=riddle_select, outputs=scenario_display)

    # Chat Interface
    chatbot = gr.ChatInterface(
        respond,
        type="messages",
        additional_inputs=[
            gr.Textbox(value="", visible=False), # Dummy für System Message (wir nutzen den generierten)
            gr.Slider(minimum=1, maximum=512, value=100, step=1, label="Max new tokens"),
            gr.Slider(minimum=0.1, maximum=2.0, value=0.2, step=0.1, label="Temperature (Niedrig für Logik)"),
            gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
            riddle_select # Das ausgewählte Rätsel wird an die respond-Funktion übergeben
        ],
    )

if __name__ == "__main__":
    demo.launch()