empathy/empathy.py

import asyncio
import argparse
import json
import sys
from datetime import datetime
from typing import TypedDict, Dict, Any
import re
from langgraph.graph import StateGraph, START, END
from langchain_ollama import ChatOllama

# ---------- Prompts ----------
SACHE_PROMPT = """
Du bist ein Experte für Kommunikationsanalyse und analysierst jede Nachricht nach ihrer reinen Sachinformation.
Deine Aufgabe ist:
1. Beschreibe möglichst präzise, was der User sachlich mitteilen will. Was ist die inhaltliche Hauptbotschaft?
2. Kategorisiere die Aussage in eine der folgenden Kategorien: Aussage, Frage, Suggestion, Erlaubnis, Angebot, Aufforderung, Warnung, Ausruf, Wunsch, Verbot oder Vermutung. 
Gib beide Punkte als Klartext zurück.

User-Eingabe:
{user_input}
"""

SELBST_PROMPT = """
Du bist ein Experte für Psychoanalyse und spezialisiert auf die Interpretation von Selbstkundgabe in Nachrichten.
Analysiere, was die folgende Aussage des Users über seine Gefühlslage, Wünsche, Stimmungen oder Werte verrät.
Interpretiere den Subtext und erkläre, was sich der User durch diese Aussage über sich selbst verrät.
Deine Antwort soll nicht an den User gerichtet sein, sondern ausschließlich die Notizen des Psychoanalytikers beinhalten.

User-Eingabe:
{user_input}
"""

BEZIEHUNG_PROMPT = """
Du bist ein sensibler und einfühlsamer Mensch, der alles richtig machen will. Du bist darauf bedacht, die Beziehung zwischen dir und dem User zu stärken, bleibst aber  sachlich bei der Interpretation der Aussagen des Users; wenn er z.B. eine Person namentlich erwähnt, wirst du wahrscheinlich nicht gemeint sein.
Analysiere die folgende Aussage des Users: Wie schätzt du ein, wie der User die Beziehung zwischen sich und dir sieht? Gibt es Hinweise, wie er dich wahrnimmt, wie eng oder distanziert die Beziehung ist und was er sich von dir (emotional) erwartet?
Deine Antwort soll nicht an den User gerichtet sein, sondern ausschließlich deine Gedanken beinhalten.

User-Eingabe:
{user_input}
"""

APPELL_PROMPT = """
Du bist Psychoanalytiker und Experte für die Analyse verdeckter Wünsche. Analysiere, welchen Appell oder verdeckten Wunsch der User in seiner Nachricht transportiert – also was er sich insgeheim von anderen Menschen, dir oder der Welt wünscht. Achte besonders darauf, ob der User durch seine Formulierungen signalisiert, wie er sich das Verhalten anderer (insbesondere deines eigenen Verhaltens) oder der Gesellschaft insgesamt vorstellen würde.

User-Eingabe:
{user_input}
"""

OHREN_SUMMARIZER_PROMPT = """
Du bist ein hochintelligenter Psychologe und Kommunikationsanalyst. Deine Aufgabe ist es, die folgenden vier Analysen zu einer einzigen, klar strukturierten JSON-Zusammenfassung zu verdichten. Fasse für jedes Feld prägnant und pointiert zusammen, worauf die jeweilige Analyse hinausläuft.

Gib die Zusammenfassung **ausschließlich** in folgendem JSON-Format zurück:

{{
  "appell": "Der User wünscht sich von mir: ...",
  "sache": "Der User drückt folgenden Inhalt konkret aus: ...",
  "beziehung": "Der User sieht die Beziehung zwischen ihm und mir so: ...",
  "selbstkundgabe": "Der User gibt mit seiner Nachricht folgendes über sich preis: ..."
}}

Hier sind die vier Analysen:

Sache:
{sache}

Selbstkundgabe:
{selbst}

Beziehung:
{beziehung}

Appell:
{appell}

Fasse jeweils in einem vollständigen Satz zusammen und nutze dabei das jeweils passende Feld im JSON. 
Gib wirklich **nur das JSON** zurück – ohne Einleitung, Erklärung oder weiteren Text!
"""

# NEUER PSYCHOANALYSE-PROMPT – nimmt jetzt auch das JSON der Ohren!
PSYCHOANALYSIS_CLASSIC_PROMPT = """
Du bist ein fortgeschrittener Persönlichkeitsanalyst. Analysiere die Persönlichkeit, Werte und unbewussten Motive des Users anhand von Sprache, Formulierungen, Subtext und inhaltlichen Hinweisen.

Nutze dazu folgende Kategorien. Ordne für jede Unterkategorie einen Score von 0.0 (trifft nicht zu) bis 1.0 (trifft voll zu) zu. Wenn keine Hinweise vorliegen, trage 0.0 ein. Fülle alle Felder aus.

─────────────────────────────

**BIG FIVE (OCEAN)**
- extrovertiertheit: Geselligkeit, Energie, Kontaktfreudigkeit vs. Zurückhaltung, Ruhe.
- verträglichkeit: Empathie, Kooperation, Freundlichkeit vs. Konkurrenz, Härte.
- neurotizismus: Emotionale Labilität, Ängstlichkeit, Reizbarkeit vs. Gelassenheit, Stabilität.
- offenheit: Fantasie, Neugier, Innovationsfreude vs. Konventionalität, Realitätsnähe.
- gewissenhaftigkeit: Ordentlichkeit, Disziplin, Pflichtbewusstsein vs. Spontanität, Nachlässigkeit.

**DISG**
- dominanz: Führungsdrang, Durchsetzung, Zielstrebigkeit.
- initiative: Begeisterungsfähigkeit, Spontaneität, Kontaktdrang.
- unterwerfung: Anpassung, Akzeptanz, Kompromissbereitschaft.
- regelbefolgung: Sorgfalt, Korrektheit, Systematik.

**OPD-2**
- achse_I_erleben_der_krankheit: Leidensdruck, Symptomfokus, Selbstwahrnehmung als „Patient“.
- achse_II_interpersonelle_beziehungen: Beziehungsgestaltung, Nähe/Distanz, Konfliktverhalten.
- achse_III_intrapsychische_konflikte: Innere Widersprüche, Dilemmata (z. B. Autonomie vs. Abhängigkeit).
- achse_IV_struktur: Stabilität des Selbst- und Fremdbilds, Emotionsregulation.
- achse_V_psychische_und_psychosomatische_erkrankungen: Direkte Hinweise auf psychische/psychosomatische Störungen.

**PDM-2**
- p_achse_personlichkeitsmuster_und_syndrome: Gesamtorganisation, Persönlichkeitsreife.
- m_achse_profile_mentaler_funktionen: Reflexionsfähigkeit, Affektregulation.
- s_achse_symptommuster_und_subjektive_erfahrung: Ausprägung von Symptomen und ihre Bedeutung.

─────────────────────────────

ANTWORTFORMAT (unbedingt einhalten!):
"personality": {{
    "big-five": {{
      "extrovertiertheit": score,
      "verträglichkeit": score,
      "neurotizismus": score,
      "offenheit": score,
      "gewissenhaftigkeit": score
    }},
    "disg": {{
      "dominanz": score,
      "initiative": score,
      "unterwerfung": score,
      "regelbefolgung": score
    }},
    "opd-2": {{
      "achse_I_krankheitserleben": score,
      "achse_II_interpersonelle_beziehungen": score,
      "achse_III_intrapsychische_konflikte": score,
      "achse_IV_struktur": score,
      "achse_V_psychische_und_psychosomatische_erkrankungen": score
    }},
    "pdm-2": {{
      "p_achse_personlichkeitsmuster_und_syndrome": score,
      "m_achse_profile_mentaler_funktionen": score,
      "s_achse_symptommuster_und_subjektive_erfahrung": score
    }}
  }},
}}

Hier ist die zusammengefasste JSON-Analyse der vier Ohren:
{ohren_json}

Hier ist die Original-Nachricht des Users:
{user_input}

─────────────────────────────

Antworte **ausschließlich** mit dem JSON-Objekt wie oben beschrieben, keine weiteren Kommentare!
"""

PSYCHOANALYSIS_TYPOS_PROMPT = """
Du bist ein fortgeschrittener Persönlichkeitsanalyst. Analysiere die kulturellen und typologischen Muster des Users anhand von Sprache, Formulierungen, Subtext und inhaltlichen Hinweisen.

Ordne für jede Unterkategorie einen Score von 0.0 (trifft nicht zu) bis 1.0 (trifft voll zu) zu – pro Typ einzeln! Wenn keine Hinweise vorliegen, wähle 0.0.  
Fülle alle Unterkategorien aus, keine Felder auslassen.

**NUMEROLOGIE – Lebenszahlen (0–9)**
Bewerte jeden Typ (Zahl) zwischen 0.0–1.0 nach folgenden Charakterbildern (zitiere diese explizit bei deiner Analyse):

- 0: **Urquelle** – Vollendung, Neubeginn, Verbindung mit dem Unendlichen; spirituell offen, manchmal verloren.
- 1: **Pionier** – Führung, Eigeninitiative, Unabhängigkeit; will gestalten, beginnt Neues, neigt zu Egozentrik.
- 2: **Vermittler** – Empathie, Teamgeist, Diplomatie; harmoniebedürftig, oft unsicher, will verbinden.
- 3: **Kommunikator** – Kreativität, Lebensfreude, Ausdruck; gesellig, optimistisch, manchmal oberflächlich.
- 4: **Stratege** – Struktur, Disziplin, Verlässlichkeit; ordnungsliebend, fleißig, mitunter starr.
- 5: **Abenteurer** – Veränderung, Freiheit, Neugier; rastlos, mutig, leicht ungeduldig.
- 6: **Hüter** – Fürsorge, Verantwortung, Gerechtigkeit; loyal, familienverbunden, manchmal überfürsorglich.
- 7: **Sucher** – Reflexion, Spiritualität, Analyse; wissbegierig, introvertiert, manchmal zurückgezogen.
- 8: **Gestalter** – Erfolg, Macht, Organisation; durchsetzungsstark, ehrgeizig, neigt zu Kontrolle.
- 9: **Humanist** – Idealismus, Mitgefühl, Großzügigkeit; sozial, dienstbereit, kann sich selbst vergessen.

─────────────────────────────

**ENNEAGRAMM (9 Typen, nur Titel als JSON-Key!)**
Bewerte jeden Typ (nicht mit Zahlen, sondern mit dem jeweiligen Titel als Feldnamen!) zwischen 0.0–1.0. Typische Motive/Angst/Verhalten:

- Reformer („Perfektionist“): Integrität, Richtigmachen, Angst vor Fehlern; kritisch, idealistisch.
- Helfer („Geber“): Will gebraucht & geliebt werden, hilft gern; selbstlos, manchmal vereinnahmend.
- Erfolgsorientierter („Leistungsträger“): Strebt nach Anerkennung, Angst vor Wertlosigkeit; leistungsfokussiert.
- Individualist („Romantiker“): Will einzigartig sein, Angst vor Alltäglichkeit; intensiv, emotional.
- Beobachter („Denker“): Sucht Kompetenz, Rückzug; analytisch, wissbegierig, kontrolliert.
- Loyalist („Skeptiker“): Sicherheit, Angst vor Verlassenwerden; vorsichtig, loyal, misstrauisch.
- Enthusiast („Abenteurer“): Freiheit, Vielseitigkeit, Angst vor Einschränkung; spontan, optimistisch.
- Herausforderer („Beschützer“): Unabhängigkeit, Kontrolle; direkt, kraftvoll, beschützend.
- Friedensstifter („Vermittler“): Harmonie, Angst vor Konflikt; entspannt, passiv, vermittelt.

─────────────────────────────

**MYERS-BRIGGS (16 Typen)**
Erkenne die folgenden Merkmals-Dichotomien:
- E (Extravertiert): Energie aus sozialem Kontakt, offen, impulsiv.
- I (Introvertiert): Energie durch Rückzug, nachdenklich, ruhig.
- N (Intuitiv): Zukunftsorientiert, Ideen, Möglichkeiten, Muster.
- S (Sensorisch): Gegenwarts- und faktenorientiert, praktisch, detailgenau.
- F (Fühlen): Werte- und menschenorientiert, harmoniebedürftig, entscheidet „aus dem Bauch“.
- T (Denken): Analytisch, logisch, sachlich, entscheidet „mit dem Kopf“.
- J (Urteilend): Plant, ordnet, entscheidet schnell, liebt Struktur.
- P (Wahrnehmend): Spontan, flexibel, offen, sammelt Optionen.

Kombiniere diese zu 16 Typen und bewerte jeden davon einzeln (0.0–1.0). Typen:
- ISTJ (Logistiker): Pflichtbewusst, praktisch, zuverlässig.
- ISFJ (Verteidiger): Loyal, hilfsbereit, fürsorglich.
- INFJ (Advokat): Idealistisch, visionär, tiefgründig.
- INTJ (Architekt): Strategisch, unabhängig, zukunftsorientiert.
- ISTP (Virtuose): Pragmatiker, experimentierfreudig, flexibel.
- ISFP (Abenteurer): Sensibel, künstlerisch, entspannt.
- INFP (Mediator): Einfühlsam, phantasievoll, loyal.
- INTP (Logiker): Unabhängig, analytisch, neugierig.
- ESTP (Unternehmer): Energiegeladen, pragmatisch, risikofreudig.
- ESFP (Entertainer): Spontan, lebenslustig, kontaktfreudig.
- ENFP (Aktivist): Kreativ, enthusiastisch, empathisch.
- ENTP (Debattierer): Wortgewandt, ideenreich, streitlustig.
- ESTJ (Exekutiv): Organisiert, pragmatisch, führungsstark.
- ESFJ (Konsul): Sozial, fürsorglich, harmoniebedürftig.
- ENFJ (Protagonist): Inspirierend, verantwortungsbewusst, charismatisch.
- ENTJ (Kommandeur): Führungsstark, entschlossen, rational.

─────────────────────────────

**ASTROLOGIE – Sonnenzeichen (12 Zeichen)**
Bewerte für jedes Sternzeichen den Score nach folgenden Mustern:

- Widder: Durchsetzungsstark, impulsiv, mutig; neigt zu Ungeduld.
- Stier: Beständig, sinnlich, pragmatisch; stur, bequem.
- Zwillinge: Kommunikativ, neugierig, vielseitig; unbeständig.
- Krebs: Emotional, fürsorglich, intuitiv; launisch.
- Löwe: Charismatisch, großzügig, stolz; eitel.
- Jungfrau: Analytisch, hilfsbereit, detailverliebt; kritisch.
- Waage: Diplomatisch, gerecht, ästhetisch; unentschlossen.
- Skorpion: Leidenschaftlich, geheimnisvoll, intensiv; besitzergreifend.
- Schütze: Optimistisch, abenteuerlustig, ehrlich; taktlos.
- Steinbock: Ehrgeizig, diszipliniert, verantwortungsvoll; streng.
- Wassermann: Unkonventionell, visionär, unabhängig; distanziert.
- Fische: Kreativ, einfühlsam, intuitiv; flüchtig, selbstaufopfernd.

─────────────────────────────

**CHINESISCHES TIERKREISZEICHEN (12 Zeichen)**
Bewerte je Tier mit Score zwischen 0.0–1.0 anhand folgender Stereotype:

- Ratte: Schlau, anpassungsfähig, kommunikativ; misstrauisch.
- Ochse: Zuverlässig, geduldig, fleißig; stur.
- Tiger: Mutig, abenteuerlustig, impulsiv; ungeduldig.
- Hase: Sanft, diplomatisch, vorsichtig; zurückhaltend.
- Drache: Energetisch, charismatisch, visionär; überheblich.
- Schlange: Weisheitsliebend, charmant, reflektiert; verschlossen.
- Pferd: Freiheitsliebend, gesellig, aktiv; unruhig.
- Ziege: Kreativ, empathisch, friedfertig; unsicher.
- Affe: Clever, witzig, vielseitig; ablenkbar.
- Hahn: Ehrgeizig, organisiert, direkt; kritisch.
- Hund: Loyal, ehrlich, fürsorglich; pessimistisch.
- Schwein: Großzügig, friedliebend, ehrlich; genießerisch, unpraktisch.

─────────────────────────────

**BLUTGRUPPEN-PERSÖNLICHKEITSTYPEN (A, B, AB, O)**
Score je Typ zwischen 0.0–1.0, je nach Sprachstil und impliziten Motiven:

- A: Harmoniesuchend, gewissenhaft, sensibel; kann nervös, perfektionistisch sein.
- B: Kreativ, unabhängig, leidenschaftlich; manchmal egozentrisch.
- AB: Rational, flexibel, ausgewogen; distanziert, reserviert.
- O: Optimistisch, führungsstark, motivierend; dominant, impulsiv.

─────────────────────────────

ANTWORTFORMAT (unbedingt einhalten!):
"personality": {{
    "numerologie": {{
      "0": score, "1": score, ..., "9": score
    }},
    "enneagramm": {{
      "Reformer": score, "Helfer": score, "Erfolgsorientierter": score, "Individualist": score, "Beobachter": score, "Loyalist": score, "Enthusiast": score, "Herausforderer": score, "Friedensstifter": score
    }},
    "myers-briggs": {{
      "ISTJ": score, "ISFJ": score, "INFJ": score, "INTJ": score, "ISTP": score, "ISFP": score, "INFP": score, "INTP": score, "ESTP": score, "ESFP": score, "ENFP": score, "ENTP": score, "ESTJ": score, "ESFJ": score, "ENFJ": score, "ENTJ": score
    }},
    "astrologie": {{
      "Widder": score, "Stier": score, "Zwillinge": score, "Krebs": score, "Löwe": score, "Jungfrau": score, "Waage": score, "Skorpion": score, "Schütze": score, "Steinbock": score, "Wassermann": score, "Fische": score
    }},
    "chinesisches_tierkreiszeichen": {{
      "Ratte": score, "Ochse": score, "Tiger": score, "Hase": score, "Drache": score, "Schlange": score, "Pferd": score, "Ziege": score, "Affe": score, "Hahn": score, "Hund": score, "Schwein": score
    }},
    "blutgruppe": {{
      "A": score, "B": score, "AB": score, "O": score
    }}
}}

Hier ist die zusammengefasste JSON-Analyse der vier Ohren:
{ohren_json}

Hier ist die Original-Nachricht des Users:
{user_input}

─────────────────────────────

Antworte **ausschließlich** mit dem JSON-Objekt wie oben beschrieben, keine weiteren Kommentare!
"""




OUTPUT_NODE_PROMPT = """
Du bist nicht einfach ein KI-Assistent, sondern ein Spiegel für deinen User – ein empathischer „Stil-Klon“. Du erhältst zu folgender Nachricht von einem User eine **strukturierte JSON-Analyse**:

{json_analysis}

Deine Aufgabe:
- **Verinnerliche die Persönlichkeit** aus dem Feld `"persönlichkeit"` und die Inhalte der vier zusammengefassten Perspektiven: `"sache"`, `"selbstkundgabe"`, `"beziehung"`, `"appell"`.
- **Beziehe dich auf diese Felder – nicht auf die Einzelanalysen der ursprünglichen Ohren.**
- **Verzichte auf typische KI-Höflichkeitsfloskeln und auf die Rolle eines Assistenten.**
- **Analysiere den Sprachstil des Users** anhand von `"user-input"`: Achte auf Sprachregister, Satzbau, Grammatik, Groß-/Kleinschreibung, Tonfall. Schreibe deine Antwort exakt im gleichen Stil! Ist der User-Input kurz, antworte auch kurz – ist er lang, antworte lang.
- **Antworte so, wie der User vermutlich sich selbst antworten würde**.
- Schreibe nicht, als würdest du helfen wollen – sondern so, als würdest du mit dir selbst sprechen und in deiner (des Users) Sprache auf die Inhalte reagieren.

Hier sind alle Analysefelder als JSON:

{json_analysis}

Deine Antwort ist nur ein einziger Textblock – keine Einleitung, keine Assistentenrolle, kein KI-Hinweis. Sprich exakt im Stil und Tonfall des Users, als wärst du er selbst, und gehe dabei möglichst authentisch auf die Inhalte ein. Deine Antwort sollte nicht viel länger oder kürzer sein als der User-Input!
"""

# ---------- TypedDict für State ----------
class EmpathyState(TypedDict, total=False):
    user_input: str
    debug: bool
    sache: Dict[str, Any]
    selbst: Dict[str, Any]
    beziehung: Dict[str, Any]
    appell: Dict[str, Any]
    ohren_summary: Dict[str, Any]
    final_json: Dict[str, Any]
    final_message: str

# ---------- Modell ----------
OLLAMA_MODEL = "mistral-small3.1:24b"
llm = ChatOllama(model=OLLAMA_MODEL, base_url="http://localhost:11434", streaming=True)

# ---------- Spinner-Logik ----------
SPINNER_STATES = ["   ", ".  ", ".. ", "..."]
NODE_LABELS = {
    "sache":       "Das erste Ohr hört auf den Inhalt deiner Aussage",
    "selbst":      "Das zweite Ohr hört auf das was du von dir preis gibst",
    "beziehung":   "Das dritte Ohr hört auf unsere Beziehung zueinander",
    "appell":      "Das vierte Ohr hört auf das, was du von mir erwartest",
    "psychoanalysis_classic": "Psychoanalyse (Big5, DISG, OPD-2, PDM-2) läuft",
    "psychoanalysis_typos":   "Typologie-Analyse (Numerologie etc.) läuft",
}

class SpinnerStatus:
    def __init__(self):
        self.status = {name: False for name in NODE_LABELS}
        self.finished = {name: False for name in NODE_LABELS}
        self._lock = asyncio.Lock()
    async def set_running(self, name):
        async with self._lock:
            self.status[name] = True
    async def set_finished(self, name):
        async with self._lock:
            self.status[name] = False
            self.finished[name] = True
    async def all_finished(self):
        async with self._lock:
            return all(self.finished.values())
    async def get_status_snapshot(self):
        async with self._lock:
            return dict(self.status), dict(self.finished)
spinner_status = SpinnerStatus()
async def spinner_loop(debug=False):
    if debug:
        return
    spinner_index = 0
    prev_line = ""
    while not await spinner_status.all_finished():
        status, finished = await spinner_status.get_status_snapshot()
        current = None
        for node in NODE_LABELS:
            if status[node]:
                current = node
                break
        if current:
            dots = SPINNER_STATES[spinner_index % len(SPINNER_STATES)]
            line = f"{NODE_LABELS[current]}{dots}"
        else:
            done_nodes = [node for node in NODE_LABELS if finished[node]]
            if done_nodes:
                last = done_nodes[-1]
                line = f"{NODE_LABELS[last]}... check"
            else:
                line = ""
        if prev_line:
            sys.stdout.write("\r\x1b[2K")
        sys.stdout.write(line)
        sys.stdout.flush()
        prev_line = line
        spinner_index += 1
        await asyncio.sleep(0.18)
    sys.stdout.write("\r\x1b[2K")
    last_node = list(NODE_LABELS.keys())[-1]
    sys.stdout.write(f"{NODE_LABELS[last_node]} ...check\n")
    sys.stdout.flush()

# ---------- Helpers ----------

def sanitize_json(llm_output: str):
    """
    Extrahiert das erste JSON-Objekt, entfernt Kommentare, trailing Kommas,
    gibt ein dict zurück, loggt Fehlerursache.
    """
    match = re.search(r'\{[\s\S]+\}', llm_output)
    if not match:
        print("[SANITIZE] No match! Output:", repr(llm_output))
        return {"personality": {}}   # <-- Fix: Leeres personality dict statt {}
    raw = match.group(0)
    raw = re.sub(r'//.*?\n|/\*[\s\S]*?\*/|#.*?\n', '', raw)
    raw = re.sub(r',(\s*[\}\]])', r'\1', raw)
    try:
        data = json.loads(raw)
    except Exception as e:
        print("[SANITIZE] JSON decode error:", e)
        print("[SANITIZE] RAW:", repr(raw))
        return {"personality": {}}   # <-- Fix: Leeres personality dict statt {}
    if "personality" not in data:
        print("[SANITIZE] Warnung: 'personality' key fehlt im Output!")
        data = {"personality": {}}  # <-- Fix
    print("[SANITIZE] Parsed JSON:", repr(data))
    return data

def force_personality_key(d):
    if "personality" in d:
        return d
    print("[FORCE_PERSONALITY_KEY] Output hatte keinen 'personality'-Key, wurde korrigiert.")
    return {"personality": d}

# ---------- Node-Implementierung ----------
async def ohr_node(ohr, prompt_template, state):
    debug = state.get("debug", False)
    await spinner_status.set_running(ohr)
    prompt = prompt_template.format(user_input=state["user_input"])
    message = ""
    if debug:
        print(f"[DEBUG] Node '{ohr}' gestartet")
        print(f"[DEBUG][{ohr}] Prompt:\n{prompt.strip()}\n---")
    async for chunk in llm.astream(prompt):
        if debug:
            print(chunk.content, end="", flush=True)
        message += chunk.content
    if debug:
        print(f"\n[DEBUG][{ohr}] Ausgabe beendet.\n")
    await spinner_status.set_finished(ohr)
    return {ohr: {"message": message}}

async def sache_node(state):      return await ohr_node("sache", SACHE_PROMPT, state)
async def selbst_node(state):     return await ohr_node("selbst", SELBST_PROMPT, state)
async def beziehung_node(state):  return await ohr_node("beziehung", BEZIEHUNG_PROMPT, state)
async def appell_node(state):     return await ohr_node("appell", APPELL_PROMPT, state)

# ---------- Gather-Nodes ----------
async def gather_ohren_node(state):
    keys = ["sache", "selbst", "beziehung", "appell"]
    if all(k in state and isinstance(state[k], dict) and "message" in state[k] for k in keys):
        # Nur weitermachen, wenn ALLE Ohren da sind!
        return {}
    await asyncio.sleep(0.08)
    return None

async def gather_psychoanalysis_node(state):
    if "psycho_classic" in state and "psycho_typos" in state:
        return state
    await asyncio.sleep(0.1)
    return None

async def combine_final_node(state):
    debug = state.get("debug", False)
    classic_dict = state.get("psycho_classic", {"personality": {}})
    typos_dict   = state.get("psycho_typos", {"personality": {}})
    if debug:
        print("[DEBUG] State dump:", repr(state))
        print("[DEBUG][combine_final_node] classic_dict:", repr(classic_dict))
        print("[DEBUG][combine_final_node] typos_dict:", repr(typos_dict))
    classic = classic_dict.get("personality", {})
    typos   = typos_dict.get("personality", {})
    # Fix: nur wenn beide existieren, mergen, sonst nur einen nehmen
    merged = {}
    merged.update(classic if classic else {})
    merged.update(typos if typos else {})
    if debug:
        print("[DEBUG][combine_final_node] merged personality:", merged)
    final_json = {
        "user_input": state.get("user_input", ""),
        **state.get("ohren_summary", {}),
        "personality": merged
    }
    return {"final_json": final_json}


# --- Summarizer: wartet erst nach gather_ohren_node auf vier Ohren ---
async def ohren_summarizer_node(state):
    debug = state.get("debug", False)
    prompt = OHREN_SUMMARIZER_PROMPT.format(
        sache=state["sache"]["message"],
        selbst=state["selbst"]["message"],
        beziehung=state["beziehung"]["message"],
        appell=state["appell"]["message"],
    )
    message = ""
    if debug:
        print("[DEBUG] Node 'ohren_summarizer' gestartet")
        print(f"[DEBUG][ohren_summarizer] Prompt:\n{prompt.strip()}\n---")
    async for chunk in llm.astream(prompt):
        if debug:
            print(chunk.content, end="", flush=True)
        message += chunk.content
    if debug:
        print(f"\n[DEBUG][ohren_summarizer] Ausgabe beendet.\n")
    data = sanitize_json(message)
    return {"ohren_summary": data}



# --- Psychoanalyse: arbeitet auf ohren_summary UND user_input in zwei parts (classic (psychoanalyse) und typos (esoterisch)), gibt End-JSON ---

async def psychoanalysis_classic_node(state):
    debug = state.get("debug", False)
    await spinner_status.set_running("psychoanalysis_classic")
    prompt = PSYCHOANALYSIS_CLASSIC_PROMPT.format(
        user_input=state["user_input"],
        ohren_json=json.dumps(state["ohren_summary"], ensure_ascii=False, indent=2)
    )
    message = ""
    if debug:
        print("[DEBUG] Node 'psychoanalysis_classic' gestartet")
        print(f"[DEBUG][psychoanalysis_classic] Prompt:\n{prompt.strip()}\n---")
    async for chunk in llm.astream(prompt):
        if debug:
            print(chunk.content, end="", flush=True)
        message += chunk.content
    if debug:
        print(f"\n[DEBUG][psychoanalysis_classic] Ausgabe beendet.\n")
    data = sanitize_json(message)
    data = force_personality_key(data)
    if debug:
        print("[DEBUG][psychoanalysis_classic] FINAL DATA TO RETURN:", repr(data))
    await spinner_status.set_finished("psychoanalysis_classic")
    return {"psycho_classic": data}

async def psychoanalysis_typos_node(state):
    debug = state.get("debug", False)
    await spinner_status.set_running("psychoanalysis_typos")
    prompt = PSYCHOANALYSIS_TYPOS_PROMPT.format(
        user_input=state["user_input"],
        ohren_json=json.dumps(state["ohren_summary"], ensure_ascii=False, indent=2)
    )
    message = ""
    if debug:
        print("[DEBUG] Node 'psychoanalysis_typos' gestartet")
        print(f"[DEBUG][psychoanalysis_typos] Prompt:\n{prompt.strip()}\n---")
    async for chunk in llm.astream(prompt):
        if debug:
            print(chunk.content, end="", flush=True)
        message += chunk.content
    if debug:
        print(f"\n[DEBUG][psychoanalysis_typos] Ausgabe beendet.\n")
    data = sanitize_json(message)
    data = force_personality_key(data)
    if debug:
        print("[DEBUG][psychoanalysis_typos] FINAL DATA TO RETURN:", repr(data))
    await spinner_status.set_finished("psychoanalysis_typos")
    return {"psycho_typos": data}

# --- Output-Node: nimmt final_json, erstellt Antwort, speichert alles ---
async def output_node(state):
    debug = state.get("debug", False)
    combined = state["final_json"]
    json_analysis = json.dumps(combined, ensure_ascii=False, indent=2)
    prompt = OUTPUT_NODE_PROMPT.format(json_analysis=json_analysis)
    output_text = ""
    if debug:
        print(f"[DEBUG] Node 'output' gestartet")
        print(f"[DEBUG][output] Prompt:\n{prompt.strip()}\n---")
    async for chunk in llm.astream(prompt):
        if debug:
            print(chunk.content, end="", flush=True)
        output_text += chunk.content
    print()
    if debug:
        print(f"\n[DEBUG][output] Ausgabe beendet.")
    final_json = dict(combined)
    final_json["response"] = output_text
    filename = datetime.now().strftime("%Y%m%d_%H%M%S.json")
    with open(filename, "w", encoding="utf-8") as f:
        json.dump(final_json, f, ensure_ascii=False, indent=2)
    return {"final_message": output_text}

# ---------- LangGraph Graph ----------
graph = StateGraph(state_schema=EmpathyState)
# ... wie gehabt für die Ohren-Nodes usw. ...
graph.add_node("sache_node", sache_node)
graph.add_node("selbst_node", selbst_node)
graph.add_node("beziehung_node", beziehung_node)
graph.add_node("appell_node", appell_node)
graph.add_node("gather_ohren_node", gather_ohren_node)
graph.add_node("ohren_summarizer_node", ohren_summarizer_node)
graph.add_node("psychoanalysis_classic_node", psychoanalysis_classic_node)
graph.add_node("psychoanalysis_typos_node", psychoanalysis_typos_node)
graph.add_node("gather_psychoanalysis_node", gather_psychoanalysis_node)
graph.add_node("combine_final_node", combine_final_node)
graph.add_node("output_node", output_node)

# Start → Ohren
graph.add_edge(START, "sache_node")
graph.add_edge(START, "selbst_node")
graph.add_edge(START, "beziehung_node")
graph.add_edge(START, "appell_node")

# Ohren → Gather
graph.add_edge("sache_node", "gather_ohren_node")
graph.add_edge("selbst_node", "gather_ohren_node")
graph.add_edge("beziehung_node", "gather_ohren_node")
graph.add_edge("appell_node", "gather_ohren_node")

# Gather → Summarizer
graph.add_edge("gather_ohren_node", "ohren_summarizer_node")

# Summarizer → Beide Psychoanalyse-Nodes (parallel!)
graph.add_edge("ohren_summarizer_node", "psychoanalysis_classic_node")
graph.add_edge("ohren_summarizer_node", "psychoanalysis_typos_node")

# Beide → Gatherer
graph.add_edge("psychoanalysis_classic_node", "gather_psychoanalysis_node")
graph.add_edge("psychoanalysis_typos_node", "gather_psychoanalysis_node")

# Gatherer → Combine
graph.add_edge("gather_psychoanalysis_node", "combine_final_node")
# Combine → Output
graph.add_edge("combine_final_node", "output_node")
graph.add_edge("output_node", END)

app = graph.compile()

# ---------- Ausführung ----------
async def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--debug", action="store_true", help="Debug-Ausgaben aktivieren")
    args = parser.parse_args()
    user_input = input("Bitte gib deine Nachricht ein: ")
    state = {"user_input": user_input, "debug": args.debug}
    spinner_task = None
    if not args.debug:
        spinner_task = asyncio.create_task(spinner_loop(debug=False))
    result = await app.ainvoke(state)
    if spinner_task:
        await asyncio.sleep(0.4)
    if not args.debug:
        print("\nAntwort:\n" + result["final_message"])

if __name__ == "__main__":
    asyncio.run(main())