import gradio as gr
from openai import OpenAI
import os
import json

# ==========================================
# 1. KONFIGURATION & AUFGABEN-PROFILE
# ==========================================

MODERATOR_MODEL = "openai/gpt-oss-120b"
MODERATOR_PROVIDER = "groq"  # 435 tok/s — wird oft aufgerufen, Geschwindigkeit summiert sich

# Jedes Profil definiert, wie Experten denken und was am Ende rauskommt
TASK_PROFILES = {
    "CODE": {
        "label": "💻 Code / Technisch",
        "expert_focuses": [
            "Architektur, Algorithmus & Gesamtstruktur des Codes",
            "Edge-Cases, Fehlerbehandlung & Sicherheit",
            "Saubere, idiomatische Implementierung & Lesbarkeit",
        ],
        "final_instruction": (
            "Liefere AUSSCHLIESSLICH den fertigen, lauffähigen Code. "
            "Inline-Kommentare nur wo nötig. Kein Fließtext drumherum. "
            "Korrekter Code-Block für die jeweilige Sprache."
        ),
        "draft_label": "Aktueller Code-Entwurf",
    },
    "TEXT": {
        "label": "✍️ Text / Content",
        "expert_focuses": [
            "Argumentationsstruktur, roter Faden & Aufbau",
            "Ton, Zielgruppe, Fakten & fehlende Tiefe",
            "Formulierung, Wirkung & plattformgerechtes Format",
        ],
        "final_instruction": (
            "Liefere den fertigen, copy-paste-fähigen Text. "
            "Passe Länge, Ton und Format EXAKT an die Plattform an "
            "(z.B. LinkedIn: kurze Absätze, mobile-freundlich, sinnvolle Emojis). "
            "Keine Platzhalter, keine KI-Floskeln, kein Meta-Kommentar."
        ),
        "draft_label": "Aktueller Text-Entwurf",
    },
    "PLAN": {
        "label": "📋 Strategie / Plan",
        "expert_focuses": [
            "Gesamtstrategie, Phasen & logische Abfolge",
            "Risiken, Abhängigkeiten, KPIs & blinde Flecken",
            "Konkrete Maßnahmen, Verantwortlichkeiten & Zeitplan",
        ],
        "final_instruction": (
            "Liefere einen klaren, strukturierten Aktionsplan. "
            "Phasen, Meilensteine, Maßnahmen. Tabellen wo sinnvoll. "
            "Direkt umsetzbar, keine abstrakten Worthülsen."
        ),
        "draft_label": "Aktueller Plan-Entwurf",
    },
    "ANALYSIS": {
        "label": "🔍 Analyse / Konzept",
        "expert_focuses": [
            "Problemstruktur, Hypothesen & Analyserahmen",
            "Daten, Belege, Gegenargumente & Lücken",
            "Schlussfolgerungen, Empfehlungen & Priorisierung",
        ],
        "final_instruction": (
            "Liefere eine strukturierte Analyse: Befunde → Bewertung → Empfehlung. "
            "Faktenbasiert, präzise, keine leeren Phrasen."
        ),
        "draft_label": "Aktuelle Analyse",
    },
}

COUNCIL_MEMBERS = [
    {
        "name": "🧠 Experte I",
        "model": "openai/gpt-oss-120b",
        "tag": "STRUKTUR",
        "color": "#1a6b3c",
        "role_hint": "Erstelle oder verbessere die Grundstruktur und das logische Gerüst.",
    },
    {
        "name": "🧐 Experte II",
        "model": "deepseek-ai/DeepSeek-R1",  # Reasoning-Modell — ideal für Kritik & Fehlersuche
        "tag": "KRITIK",
        "color": "#7c3aed",
        "role_hint": (
            "Prüfe den Entwurf von [STRUKTUR] kritisch. "
            "Finde konkrete Fehler, Lücken und Schwächen — und behebe sie direkt im Entwurf."
        ),
    },
    {
        "name": "🛠️ Experte III",
        "model": "moonshotai/Kimi-K2.5",  # andere Trainings-Lineage → echte Diversität im Council
        "code_model": "Qwen/Qwen3-Coder-480B-A35B-Instruct",  # Swap für CODE-Tasks
        "tag": "UMSETZUNG",
        "color": "#b45309",
        "role_hint": (
            "Nimm den von [KRITIK] überarbeiteten Entwurf und bringe ihn zur Serienreife. "
            "Schärfe die Formulierungen, vervollständige fehlende Teile, sorge für Konsistenz."
        ),
    },
]


# ==========================================
# 2. PROMPT MANAGER
# ==========================================

class PromptManager:

    @staticmethod
    def task_detection_sys():
        return (
            "Du bist ein präziser Aufgaben-Klassifikator. Analysiere die Anfrage und antworte NUR mit einem JSON-Objekt.\n"
            "Gültige task_type-Werte: CODE, TEXT, PLAN, ANALYSIS\n"
            "Format (keine weiteren Zeichen außerhalb):\n"
            '{"task_type": "...", "core_goal": "Ein-Satz-Beschreibung des Ziels", '
            '"key_constraints": ["Constraint 1", "Constraint 2"]}'
        )

    @staticmethod
    def task_detection_user(user_prompt):
        return f"Aufgabe: {user_prompt}"

    @staticmethod
    def moderator_kickoff_sys(task_profile):
        return (
            f"Du bist Lead-Moderator eines Expertenrats. Aufgabentyp: {task_profile['label']}.\n"
            "Brief das Team in genau 4 Punkten (je 1 Satz):\n"
            "1. Konkretes Ziel\n"
            "2. Wichtigste Qualitätskriterien für das Endprodukt\n"
            "3. Größte Risiken / häufigste Fehler bei diesem Aufgabentyp\n"
            "4. Erwartetes Format des Endprodukts\n"
            "Kein Smalltalk. Direkt. Ohne Anrede."
        )

    @staticmethod
    def moderator_kickoff_user(user_prompt, task_info):
        constraints = ", ".join(task_info.get("key_constraints", [])) or "keine"
        return (
            f"Auftrag: '{user_prompt}'\n"
            f"Kernziel: {task_info.get('core_goal', '')}\n"
            f"Constraints: {constraints}"
        )

    @staticmethod
    def moderator_steering_sys():
        return (
            "Du bist Lead-Moderator. Bewerte den aktuellen Entwurf knapp und gib dann "
            "EINEN einzigen, konkreten Arbeitsauftrag für die nächste Runde.\n"
            "Format:\n"
            "STAND: [Was gut ist — 1 Satz]\n"
            "AUFTRAG: [Was als nächstes konkret zu tun ist — 1-2 Sätze, so spezifisch wie möglich]"
        )

    @staticmethod
    def moderator_steering_user(current_draft, round_num):
        return (
            f"Aktueller Entwurf nach Runde {round_num}:\n\n{current_draft}\n\n"
            "Gib Steuerungsanweisung für Runde {next_round}.".replace(
                "{next_round}", str(round_num + 1)
            )
        )

    @staticmethod
    def expert_sys(expert, task_profile, focus_area, round_num):
        tag = expert["tag"]
        return (
            f"Du bist Experte [{tag}] in einem iterativen Expertenrat.\n\n"
            f"DEIN FOKUS in dieser Runde: {focus_area}\n"
            f"DEINE ROLLE [{tag}]: {expert['role_hint']}\n\n"
            f"REGELN:\n"
            f"- Beginne mit '[{tag}] '\n"
            f"- Du lieferst den {task_profile['draft_label']} — nicht einen Kommentar darüber.\n"
            f"- Verbessere den Entwurf direkt. Kein 'Ich würde vorschlagen...'\n"
            f"- Keine Wiederholung des Auftrags, kein Meta-Kommentar am Ende.\n"
            f"- Runde {round_num}: {'Erstelle die erste Version.' if round_num == 1 and tag == 'STRUKTUR' else 'Baue auf dem bestehenden Entwurf auf.'}"
        )

    @staticmethod
    def expert_user(user_prompt, current_draft, steering_instruction, expert_tag):
        draft_block = (
            f"\n\nAKTUELLER ENTWURF (verbessere diesen direkt):\n---\n{current_draft}\n---"
            if current_draft
            else "\n\nEs gibt noch keinen Entwurf. Erstelle die erste Version."
        )
        steering_block = (
            f"\n\nMODERATOR-ANWEISUNG: {steering_instruction}"
            if steering_instruction
            else ""
        )
        return (
            f"Auftrag: '{user_prompt}'"
            f"{draft_block}"
            f"{steering_block}\n\n"
            f"Liefere jetzt den verbesserten {expert_tag}-Entwurf:"
        )

    @staticmethod
    def final_sys(task_profile):
        return (
            f"Du bist ein Output-Finisher für {task_profile['label']}-Aufgaben.\n\n"
            f"ANWEISUNG:\n{task_profile['final_instruction']}\n\n"
            "Nutze den vorliegenden Entwurf als Basis und liefere das polierte Endprodukt. "
            "Kein einleitender Satz, kein abschließender Kommentar — nur das Produkt."
        )

    @staticmethod
    def final_user(user_prompt, best_draft, task_info):
        return (
            f"Ursprünglicher Auftrag: '{user_prompt}'\n"
            f"Kernziel: {task_info.get('core_goal', '')}\n\n"
            f"Bester Entwurf aus der Diskussion:\n---\n{best_draft}\n---\n\n"
            "Erstelle das finale, sofort nutzbare Endprodukt:"
        )


# ==========================================
# 3. LLM SERVICE & UI HELPER
# ==========================================

class LLMService:
    def __init__(self):
        self.client = OpenAI(
            base_url="https://router.huggingface.co/v1",
            api_key=os.getenv("HF_TOKEN"),
        )

    @staticmethod
    def _strip_thinking(text: str) -> str:
        """Entfernt <think>...</think> Blöcke (z.B. von DeepSeek-R1) aus dem Output."""
        import re
        return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

    def ask(self, model_id, system_prompt, user_input, provider=None):
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_input},
        ]
        extra_headers = {"x-provider": provider} if provider else {}
        try:
            response = self.client.chat.completions.create(
                model=model_id,
                messages=messages,
                max_tokens=4000,
                temperature=0.4,
                stream=False,
                extra_headers=extra_headers,
            )
            raw = response.choices[0].message.content or ""
            return self._strip_thinking(raw)
        except Exception as e:
            return f"🚨 Fehler ({model_id}): {str(e)}"


class UIHelper:
    @staticmethod
    def header(title, color="#FF5A4D"):
        return (
            f"<h2 style='color:{color}; border-bottom:2px solid #FFEBE8; "
            f"padding-bottom:5px; margin-top:20px;'>{title}</h2>"
        )

    @staticmethod
    def message(label, content, color="#4241A6"):
        return f"**<span style='color:{color}; font-size:1.05em;'>{label}</span>**\n\n{content}"

    @staticmethod
    def info(content):
        return f"> ℹ️ {content}"


# ==========================================
# 4. ORCHESTRATOR
# ==========================================

class PlenumOrchestrator:
    def __init__(self):
        self.llm = LLMService()
        self.pm = PromptManager()
        self.ui = UIHelper()

    def _detect_task(self, user_prompt):
        raw = self.llm.ask(
            MODERATOR_MODEL,
            self.pm.task_detection_sys(),
            self.pm.task_detection_user(user_prompt),
            provider=MODERATOR_PROVIDER,
        )
        try:
            clean = raw.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip()
            return json.loads(clean)
        except Exception:
            return {"task_type": "TEXT", "core_goal": user_prompt, "key_constraints": []}

    def run(self, user_prompt, rounds):
        if not user_prompt.strip():
            yield [{"role": "assistant", "content": "Bitte gib einen Auftrag ein."}]
            return

        history = [{"role": "user", "content": user_prompt}]
        yield history

        # Zustand: Das ist das Herzstück des neuen Ansatzes
        current_draft = ""
        steering_instruction = ""

        # ── SCHRITT 0: AUFGABE ERKENNEN ──────────────────────────────
        history.append({"role": "assistant", "content": self.ui.header("🔍 AUFGABENANALYSE", "#6b7280")})
        yield history

        task_info = self._detect_task(user_prompt)
        task_type = task_info.get("task_type", "TEXT")
        task_profile = TASK_PROFILES.get(task_type, TASK_PROFILES["TEXT"])

        detection_text = (
            f"**Erkannter Typ:** {task_profile['label']}  \n"
            f"**Kernziel:** {task_info.get('core_goal', '—')}  \n"
            f"**Constraints:** {', '.join(task_info.get('key_constraints', ['—'])) or '—'}"
        )
        history.append({"role": "assistant", "content": detection_text})
        yield history

        # ── SCHRITT 1: KICK-OFF ───────────────────────────────────────
        history.append({"role": "assistant", "content": self.ui.header("🎤 MODERATOR: SITZUNGSERÖFFNUNG")})
        yield history

        kickoff = self.llm.ask(
            MODERATOR_MODEL,
            self.pm.moderator_kickoff_sys(task_profile),
            self.pm.moderator_kickoff_user(user_prompt, task_info),
            provider=MODERATOR_PROVIDER,
        )
        history.append({"role": "assistant", "content": self.ui.message("🎤 Moderator", kickoff, "#FF5A4D")})
        yield history

        # ── SCHRITT 2: ZYKLEN ─────────────────────────────────────────
        for r in range(int(rounds)):
            history.append({
                "role": "assistant",
                "content": self.ui.header(f"🔄 ZYKLUS {r + 1} — EXPERTENDEBATTE", "#4241A6"),
            })
            yield history

            # Moderator-Steuerung (ab Runde 2, basierend auf aktuellem Entwurf)
            if r > 0 and current_draft:
                steering_instruction = self.llm.ask(
                    MODERATOR_MODEL,
                    self.pm.moderator_steering_sys(),
                    self.pm.moderator_steering_user(current_draft, r),
                    provider=MODERATOR_PROVIDER,
                )
                history.append({
                    "role": "assistant",
                    "content": self.ui.message("🎤 Moderator (Steuerung)", steering_instruction, "#FF5A4D"),
                })
                yield history

            # Experten arbeiten SEQUENZIELL:
            # Experte I sieht aktuellen Entwurf → produziert Draft A
            # Experte II sieht Draft A → produziert Draft B (mit Korrekturen)
            # Experte III sieht Draft B → produziert Draft C (serienreif)
            # → Draft C wird zum current_draft für die nächste Runde
            round_draft = current_draft

            for idx, expert in enumerate(COUNCIL_MEMBERS):
                focus = task_profile["expert_focuses"][idx]

                # Task-aware Model Swap: Experte III nutzt Coder-Modell bei CODE-Tasks
                model_id = (
                    expert.get("code_model", expert["model"])
                    if task_type == "CODE" and expert["tag"] == "UMSETZUNG"
                    else expert["model"]
                )

                sys_msg = self.pm.expert_sys(expert, task_profile, focus, r + 1)
                usr_msg = self.pm.expert_user(user_prompt, round_draft, steering_instruction if r > 0 else "", expert["tag"])

                answer = self.llm.ask(model_id, sys_msg, usr_msg, provider=expert.get("provider"))

                # Dieser Experte liefert den Entwurf für den Nächsten
                round_draft = answer

                label = f"{expert['name']} [{expert['tag']}] — {focus}"
                history.append({"role": "assistant", "content": self.ui.message(label, answer, expert["color"])})
                yield history

            # Bester Stand dieser Runde = Output von Experte III
            current_draft = round_draft
            history.append({
                "role": "assistant",
                "content": self.ui.info(f"Zyklus {r + 1} abgeschlossen. Entwurf gesichert → Basis für nächste Runde."),
            })
            yield history

        # ── SCHRITT 3: FINALE AUSGABE ─────────────────────────────────
        history.append({"role": "assistant", "content": self.ui.header("🏆 FINALE AUSGABE")})
        yield history

        final = self.llm.ask(
            MODERATOR_MODEL,
            self.pm.final_sys(task_profile),
            self.pm.final_user(user_prompt, current_draft, task_info),
            provider=MODERATOR_PROVIDER,
        )
        history.append({"role": "assistant", "content": final})
        yield history


# Orchestrator-Instanz
orchestrator = PlenumOrchestrator()


# ==========================================
# 5. GRADIO UI
# ==========================================

v_theme = gr.themes.Soft(
    primary_hue="indigo",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
).set(
    button_primary_background_fill="#4241A6",
    button_primary_background_fill_hover="#2D2C73",
    button_primary_text_color="white",
    block_title_text_color="#FF5A4D",
    block_label_text_color="#4241A6",
    body_text_color="#1F2937",
    color_accent_soft="#FFEBE8",
)

with gr.Blocks() as demo:
    gr.HTML("""
    <div style="text-align:center; margin-bottom:2rem; margin-top:1rem;">
        <h1 style="color:#FF5A4D; font-weight:900; font-size:2.8rem; margin-bottom:0.2rem;
                   font-family:'Inter',sans-serif; letter-spacing:-0.02em;">PromptPlenum42</h1>
        <p style="color:#4B5563; font-size:1.1rem; font-family:'Inter',sans-serif;">
            AI-Driven Multi-Agent Consensus System
        </p>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=4):
            input_text = gr.Textbox(
                label="Plenumsauftrag",
                placeholder=(
                    "z.B. 'Refaktoriere dieses Python-Skript' · "
                    "'Schreibe einen LinkedIn-Post über KI' · "
                    "'Erstelle einen Go-to-Market-Plan für ein SaaS-Produkt'"
                ),
                lines=3,
            )
        with gr.Column(scale=1):
            rounds_slider = gr.Slider(
                minimum=1, maximum=5, value=1, step=1, label="Diskussionszyklen"
            )

    with gr.Row():
        start_btn = gr.Button("Sitzung starten", variant="primary", size="lg")
        clear_btn = gr.ClearButton(
            components=[input_text], value="Protokoll leeren", size="lg"
        )

    chatbot = gr.Chatbot(
        label="Sitzungsprotokoll",
        height=700,
        sanitize_html=False,
    )
    clear_btn.add(chatbot)

    input_text.submit(
        orchestrator.run, inputs=[input_text, rounds_slider], outputs=[chatbot]
    )
    start_btn.click(
        orchestrator.run, inputs=[input_text, rounds_slider], outputs=[chatbot]
    )

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