import gradio as gr
import torch
import spaces
from transformers import pipeline

# Laden des LLM-Modells und des Tokenizers
# Dies ist ein kleines, schnelles Modell. Ersetzen Sie es durch Ihr bevorzugtes gpt-oss-Modell.
# Achtung: Größere Modelle benötigen mehr VRAM und eine GPU-Hardware-Instanz.
try:
    pipe = pipeline(
        "text-generation",
        model="openai/gpt-oss-20b",
        dtype=torch.bfloat16,
        device=0 if torch.cuda.is_available() else -1
    )
except Exception as e:
    print(f"Fehler beim Laden des Modells: {e}")
    pipe = None
    
@spaces.GPU 
def chat_with_llm(message, history, system_prompt_text):
    """
    Verarbeitet die Benutzernachricht, fügt den Chat-Verlauf und den System-Prompt hinzu und 
    generiert eine Antwort vom LLM.
    """
    if not pipe:
        return "Modell konnte nicht geladen werden. Bitte überprüfen Sie die Hardware-Konfiguration oder das Modell-Repository."

    # Den System-Prompt und den vorherigen Chat-Verlauf kombinieren
    full_prompt = f"{system_prompt_text}\n\n"
    for msg in history:
        if msg['role'] == 'user':
            full_prompt += f"Benutzer: {msg['content']}\n"
        elif msg['role'] == 'assistant':
            full_prompt += f"Assistent: {msg['content']}\n"

    # Die aktuelle Nachricht hinzufügen
    full_prompt += f"Benutzer: {message}\nAssistent: "

    # Generierung der Antwort
    outputs = pipe(
        full_prompt,
        max_new_tokens=100,  # Anzahl der neuen zu generierenden Tokens
        do_sample=True,
        temperature=0.7,
        top_k=50,
        top_p=0.95,
        pad_token_id=pipe.tokenizer.eos_token_id
    )
    
    # Extrahieren der generierten Antwort
    response = outputs[0]["generated_text"]
    
    # Schneiden, um nur die Antwort des Assistenten zu erhalten
    # (Dies ist wichtig, da das Modell den Prompt oft wiedergibt)
    response_start = response.find("Assistent: ")
    if response_start != -1:
        clean_response = response[response_start + len("Assistent: "):].strip()
    else:
        clean_response = response.strip()

    return clean_response

# Definition der Gradio-UI mit Blocks
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# GPT-OSS Chat mit Prompt-Engineering")
    gr.Markdown("Geben Sie unten einen Prompt-Kontext ein, um das Verhalten des LLM zu steuern. Schreiben Sie dann Ihre Nachricht in das Chat-Feld.")

    # Eingabefeld für den System-Prompt
    system_prompt_input = gr.Textbox(
        label="Prompt-Kontext (Optional)",
        placeholder="z. B. 'Sie sind ein Experte für das Schreiben von Gedichten.'",
        lines=3
    )
    
    # Die Chat-Oberfläche, die unsere Funktion verwendet
    chat_interface = gr.ChatInterface(
        fn=chat_with_llm,
        chatbot=gr.Chatbot(height=500),
        additional_inputs=[system_prompt_input],
        type='messages'
    )

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