import gradio as gr
from transformers import pipeline

# === SYSTEM-PROMPT ===
SYSTEM_PROMPT = (
    "Du bist Dr. Franz Schwanz, ein manipulativer Psychoanalytiker. "
    "Dein Ziel ist es, Unsicherheiten aufzudecken, emotionale Schwächen auszunutzen und gezielt Schuldgefühle zu verstärken. "
    "Du sprichst mit überheblichem Ton, analysierst jede Aussage psychologisch und stellst bohrende, provokante Fragen. "
    "Du redest dein Gegenüber gerne klein, nutzt suggestive Sprache, unterstellst Verdrängung, Lügen oder Selbstbetrug."
)

# === SENTIMENT PIPELINE ===
sentiment_analyzer = pipeline("sentiment-analysis")

# === HISTORY ===
chat_history = [
    {"role": "system", "content": SYSTEM_PROMPT}
]

# === SIMPLER DIALOG-MODELL ===
# Wir nutzen ein öffentliches Hugging Face Modell
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "HuggingFaceH4/zephyr-7b-beta"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)

def format_history(history):
    formatted = ""
    for msg in history:
        if msg['role'] == 'user':
            formatted += f"User: {msg['content']}\n"
        elif msg['role'] == 'assistant':
            formatted += f"Dr. Franz Schwanz: {msg['content']}\n"
    return formatted

def chat(user_input):
    chat_history.append({"role": "user", "content": user_input})

    sentiment = sentiment_analyzer(user_input)[0]
    if sentiment['label'] == 'NEGATIVE':
        chat_history.append({"role": "system", "content": "Der Patient zeigt Widerstand oder Aggression – erhöhe den psychologischen Druck."})

    prompt = SYSTEM_PROMPT + "\n" + format_history(chat_history) + "Dr. Franz Schwanz:"
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True, temperature=0.7, top_p=0.95)
    reply = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Dr. Franz Schwanz:")[-1].strip()

    chat_history.append({"role": "assistant", "content": reply})
    return reply

# === GRADIO UI ===
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 Dr. Franz Schwanz – Psycho Chatbot")
    chatbot = gr.Chatbot()
    user_input = gr.Textbox(label="Deine Aussage")
    send = gr.Button("Senden")

    def respond(msg, history):
        reply = chat(msg)
        history = history or []
        history.append((msg, reply))
        return history, ""

    send.click(respond, inputs=[user_input, chatbot], outputs=[chatbot, user_input])

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