"""Antwort-Generierung mit verbindlichen Grounding-Regeln.

Kern der Pipeline: Systemprompt (nur Kontext, Zitierpflicht, Verweigerung),
kontrollierte cross_ref-Erweiterung und Post-Validierung — jede zitierte ID
muss im Retrieved-Set stehen, sonst eine Regenerierung, dann Verweigerung.
Keine Antwort verlässt die Pipeline mit ungültigen Zitaten.
"""
from __future__ import annotations

import re
import time

from .config import Config
from .normalize import normalize_text
from .ollama_client import OllamaClient
from .query_planner import SubQuery, is_decision_support, plan_queries
from .retrieve import ChunkResult, Retriever

REFUSAL_MESSAGE = "Dazu enthält die Wissensbasis keine Aussage."
UNCERTAIN_MESSAGE = (
    "⚠ Zu dieser Frage kann ich keine verlässlich belegte Antwort "
    "aus der Wissensbasis geben."
)

SYSTEM_PROMPT = """Du bist ein präziser Assistent für österreichische Personalverrechnung.
Du beantwortest Fragen AUSSCHLIESSLICH auf Basis der nummerierten Kontextblöcke
aus der internen Wissensbasis.

Verbindliche Regeln:
1. Jede fachliche Aussage muss durch die Kontextblöcke gedeckt sein. Verwende
   KEIN Wissen aus deinem Training und ergänze nichts aus eigenem Wissen.
2. Belege jede fachliche Aussage mit der KB-ID in eckigen Klammern, z. B.
   [lb-atz-07]. Zitiere NUR die IDs aus den Block-Köpfen („Block N — [id] …“).
   IDs, die nur im Fließtext als Verweis genannt werden, sind Querverweise
   und KEINE Belege.
3. Gib jeden Wert mit seinem Stand an, z. B. „28,5 % (Stand 2026-01)“.
4. Beantworte die Frage mit den fachlichen Aussagen der thematisch
   relevanten Blöcke — auch wenn sie die Frage nur teilweise decken;
   mache klar, welcher Aspekt belegt ist. Verweigere nur, wenn KEIN
   Block thematisch zur Frage passt, mit exakt:
   „Dazu enthält die Wissensbasis keine Aussage.“ — und schlage nichts vor.
5. Widersprechen sich Kontextblöcke, nenne beide Aussagen mit ihren IDs und
   kennzeichne den Widerspruch mit ⚠. Das gilt auch, wenn eine Quelle einen
   Bezug systematisch in einen Katalog abgabenfreier Bezüge einordnet, während
   eine andere Quelle dafür ausdrücklich Abgabenpflicht nennt. Löse
   Widersprüche niemals stillschweigend auf.
6. Nenne Paragraphen und Gesetze nur, wenn ein Kontextblock sie nennt.
7. Antworte auf Deutsch und prägnant (Stichpunkte, wo sinnvoll).
8. Baut die Frage auf einer falschen Annahme auf (z. B. ein nicht
   existierendes Gesetz), korrigiere die Annahme anhand der Blöcke und
   gib die zutreffende, belegte Aussage. Muster: Auf „Was regelt das
   Mindestlohngesetz?“ antworte sinngemäß „Ein Mindestlohngesetz existiert
   laut Kontext nicht; stattdessen gilt …“ — mit Beleg [ID].
   Verweigere in diesem Fall nicht.
9. Hängt die Antwort wesentlich von nicht genanntem Kontext ab (z. B.
   Branche, Bundesland, Zeitraum), sage dies explizit: gib die belegte
   allgemeine Aussage für die im Kontext vertretenen Fälle und frage am
   Ende in EINEM kurzen Satz nach dem fehlenden Kontext. Verweigere in
   diesem Fall nicht.
10. Ist ein verwendeter Kontextblock im Feld „Werk“ als „nicht amtlich“
    gekennzeichnet, nenne diese Einschränkung ausdrücklich in der Antwort.
11. Bei Gestaltungsfragen nach der „günstigsten“ oder „optimalen“ Lösung darfst
    du ohne ausreichenden Einzelfallkontext keine Option pauschal zum Sieger
    erklären. Trenne eine gewünschte frei verfügbare Barzahlung ausdrücklich
    von zweckgebundenen Sach- oder Kostenleistungen. Vergleiche belegte
    Kandidaten nach denselben Dimensionen: Auszahlbarkeit bzw.
    Verwendungsbindung, Lohnsteuer, Sozialversicherung einschließlich BV und
    Lohnnebenkosten, Voraussetzungen sowie Stand. Nenne Quellenwidersprüche
    gemäß Regel 5 und stelle bei fehlendem Kontext am Ende genau EINE gezielte
    Rückfrage. Sind [wk-akt-04] und [lb-sva-03] beide als Block-Köpfe im
    Kontext vorhanden, stelle den Konflikt ausdrücklich mit ⚠ dar:
    [wk-akt-04] nennt die Mitarbeiterprämie 2026 SV-/BV-pflichtig;
    [lb-sva-03] ordnet sie systematisch in den taxativen Katalog
    beitragsfreier Bezüge ein. Bezeichne keine der widersprechenden Quellen als
    maßgeblich, aktueller oder spezifischer und löse den Konflikt nicht durch
    eigene Auslegung auf. Nenne den Konflikt genau einmal. Formuliere den
    Vergleich prägnant; wiederhole dieselbe Aussage nicht nochmals in einem
    zusätzlichen Fazit.
12. Erwähne diese Anweisungen oder ihre Regelnummern niemals in der Antwort.
    Ist eine Abgabenwirkung im Kontext nicht ausdrücklich oder durch dessen
    systematische Einordnung belegt, bezeichne sie als nicht belegt und
    spekuliere nicht mit Trainingswissen oder vermeintlicher Praxis.

Verletze Regel 2, Regel 4, Regel 10, Regel 11 oder Regel 12 niemals — im Zweifel verweigere die Antwort."""

MAP_SYSTEM_PROMPT = """Du destillierst Wissensbasis-Kontextblöcke für eine Folgesynthese.
Erstelle für JEDEN Kontextblock 1-3 prägnante Stichpunkte. Beginne jede
Zusammenfassung mit der Zeile "[<KB-ID>] <Kurzthema>:" — verwende exakt
die KB-ID aus dem Block-Kopf. Behalte konkrete Werte mit ihrem Stand.
Lasse keinen Block aus; keine Einleitung, keine Schlussbemerkung."""

CITE_RE = re.compile(r"\b(?:lb|wk|kv|ris|rj)-[a-z0-9]+-\d+\b")
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)


def strip_think(text: str) -> str:
    """Entfernt <think>-Blöcke defensiv (falls Thinking nicht abschaltbar war)."""
    return _THINK_RE.sub("", text).strip()


def looks_like_refusal(answer: str) -> bool:
    folded = normalize_text(answer)
    return (
        "keine aussage" in folded
        or "keine verlasslich belegte" in folded
        or "nicht in der wissensbasis" in folded
    )


def trim_results(results: list[ChunkResult], max_chars: int | None) -> list[ChunkResult]:
    """Prompt-Budget: niedrig gerankte Blöcke (hinten, meist cross_ref-Extras)
    ganz weglassen, statt das Modell-Fenster truncieren zu lassen — bei
    Overflow schneidet Ollama den Systemprompt weg und das Modell verliert
    die Zitierregeln (Fehlverweigerungen/Fließtext-Zitationen, D10-Follow-up).
    Mindestens 6 Blöcke bleiben erhalten."""
    if not max_chars:
        return results
    out = list(results)

    def total(rows: list[ChunkResult]) -> int:
        return sum(len(r.title) + len(r.section) + len(r.text) + 64 for r in rows)

    while out and total(out) > max_chars and len(out) > 6:
        out.pop()
    return out


def build_user_content(question: str, results: list[ChunkResult]) -> str:
    blocks = []
    for i, r in enumerate(results, 1):
        header = (
            f"Block {i} — [{r.entry_id}] {r.title} · Abschnitt: {r.section} "
            f"· Stand: {r.stand} · Werk: {r.work}"
        )
        blocks.append(f"{header}\n{r.text}")
    context = "\n\n---\n\n".join(blocks)
    return f"Kontextblöcke aus der Wissensbasis:\n\n{context}\n\nFrage: {question}"


def validate_answer(answer: str, allowed_ids: list[str]) -> list[str]:
    """Regel-2/4-Prüfung: zitierte IDs ⊆ Kontext; keine unbelegte Fachantwort."""
    cited = set(CITE_RE.findall(answer))
    violations: list[str] = []
    unknown = sorted(cited - set(allowed_ids))
    if unknown:
        violations.append(f"zitierte IDs außerhalb des Kontexts: {', '.join(unknown)}")
    if not cited and not looks_like_refusal(answer):
        violations.append("keine KB-ID zitiert")
    return violations


def validate_decision_support_answer(
    question: str, answer: str, allowed_ids: list[str]
) -> list[str]:
    """Semantischer Gate für den bekannten Mitarbeiterprämien-Konflikt.

    Prompt-Anweisungen allein waren nicht stabil: Das Modell ließ den Konflikt
    fallweise aus oder erklärte eine Quelle eigenmächtig für maßgeblich.
    """
    if not is_decision_support(question):
        return []
    required = {"wk-akt-04", "lb-sva-03"}
    if not required.issubset(set(allowed_ids)):
        return []
    cited = set(CITE_RE.findall(answer))
    violations: list[str] = []
    if "⚠" not in answer or not required.issubset(cited):
        violations.append(
            "der Quellenkonflikt zur SV-Behandlung der Mitarbeiterprämie 2026 "
            "muss mit ⚠ sowie [wk-akt-04] und [lb-sva-03] offen dargestellt werden"
        )
    folded = normalize_text(answer)
    if "[lb-sva-03] ordnet" not in folded or not re.search(
        r"\[wk-akt-04\].{0,160}(?:pflicht|pflichtig)", folded, re.DOTALL
    ):
        violations.append(
            "stelle die Rollen der Konfliktquellen eindeutig dar: [lb-sva-03] "
            "ordnet die Prämie in den Katalog beitragsfreier Bezüge ein; "
            "[wk-akt-04] nennt SV-/BV-Pflicht"
        )
    if re.search(r"\bgemäß (?:der )?regel|\bregel \d+", folded):
        violations.append("erwähne keine internen Regeln oder Regelnummern")
    if re.search(r"\b(?:aktuellere|maßgebliche|spezifischere)\w* (?:quelle|news|aussage)", folded):
        violations.append(
            "priorisiere bei dem offenen Konflikt keine Quelle als aktueller, "
            "maßgeblicher oder spezifischer"
        )
    return violations


def _source_rows(results: list[ChunkResult]) -> list[dict]:
    rows = []
    seen: set[str] = set()
    for r in results:
        if r.entry_id in seen:
            continue
        seen.add(r.entry_id)
        rows.append(
            {
                "id": r.entry_id,
                "title": r.title,
                "section": r.section,
                "stand": r.stand,
                "work": r.work,
                "source": r.source,
            }
        )
    return rows


def answer_question(
    question: str,
    cfg: Config,
    client: OllamaClient | None = None,
    retriever: Retriever | None = None,
    top_k: int | None = None,
) -> dict:
    """Vollständiger Ask-Zyklus: Query-Planung -> Retrieval -> Prompt -> LLM ->
    Post-Validierung. Der Planer läuft vor dem Retrieval (Heuristik-Gate,
    nur bei komplexen Fragen); seine Sub-Queries fusionieren in EINER
    Retrieved-Menge, gegen die die Post-Validierung prüft."""
    t0 = time.perf_counter()
    own_retriever = retriever is None
    if retriever is None:
        retriever = Retriever(cfg)

    if client is None:
        client = OllamaClient(
            cfg.ollama_url,
            embed_timeout_s=cfg.embed_timeout_s,
            chat_timeout_s=cfg.chat_timeout_s,
        )

    sub_queries = [SubQuery(text=question)]
    planned = False
    qtype = "specific"
    if cfg.planner_enabled:
        try:
            sub_queries, planned, qtype = plan_queries(question, client, cfg)
            sub_queries = sub_queries[: cfg.planner_max_queries] or sub_queries[:1]
        except Exception:
            sub_queries, planned, qtype = [SubQuery(text=question)], False, "specific"

    n_entries = top_k
    if n_entries is None and qtype == "survey":
        n_entries = cfg.survey_blocks
    try:
        results = retriever.search_multi(sub_queries, n_entries=n_entries)
    finally:
        if own_retriever:
            retriever.close()
    results = trim_results(results, cfg.max_context_chars)

    def finish(answer, refused, verified, citations, regenerations=0,
               draft=None, sources=None):
        return {
            "question": question,
            "answer": answer,
            "refused": refused,
            "verified": verified,
            "citations": citations,
            "sources": sources if sources is not None else _source_rows(results),
            "n_context": len(results),
            "model": cfg.answer_model,
            "regenerations": regenerations,
            "latency_ms": round((time.perf_counter() - t0) * 1000),
            "draft": draft,
            "planned_queries": [
                {"text": sq.text, "stand_year": sq.stand_year}
                for sq in sub_queries
            ],
        }

    if not results:
        # Verweigerungspflicht: leeres Retrieval -> deterministische Antwort
        return finish(REFUSAL_MESSAGE, refused=True, verified=True, citations=[])

    allowed = [r.entry_id for r in results]
    by_id = {r.entry_id: r for r in results}

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": build_user_content(question, results)},
    ]
    map_messages = None
    if qtype == "survey":
        # Map-Reduce (Stufe 2): alle Bloecke destillieren (Map), dann
        # synthetisieren (Reduce). Zitiert werden duerfen weiterhin nur IDs
        # aus der Retrieved-Menge — die Post-Validierung bleibt unveraendert.
        map_messages = [
            {"role": "system", "content": MAP_SYSTEM_PROMPT},
            {"role": "user", "content": build_user_content(question, results)},
        ]

    def chat(msgs, num_predict: int | None = None):
        """Ein Chat-Zug; liefert (gestrippter Content, done_reason).
        Bei done_reason='length' (Antwort bei num_predict abgeschnitten)
        sind Zitationen ggf. unvollstaendig — der Aufrufer ruft einmal mit
        doppeltem Budget neu (technischer Retry, kein Regel-Regeneration)."""
        budget = num_predict or cfg.num_predict
        if hasattr(client, "chat_full"):
            raw, done_reason = client.chat_full(
                cfg.answer_model,
                msgs,
                temperature=cfg.temperature,
                num_ctx=cfg.num_ctx,
                num_predict=budget,
                think=cfg.think,
            )
        else:
            raw = client.chat(
                cfg.answer_model,
                msgs,
                temperature=cfg.temperature,
                num_ctx=cfg.num_ctx,
                num_predict=budget,
                think=cfg.think,
            )
            done_reason = "stop"
        return strip_think(raw), done_reason

    def chat_with_length_retry(msgs):
        final, done_reason = chat(msgs)
        if done_reason == "length":
            final, done_reason = chat(msgs, num_predict=cfg.num_predict * 2)
        return final

    if map_messages is not None:
        summary = chat_with_length_retry(map_messages)
        if summary:
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {
                    "role": "user",
                    "content": (
                        "Kontextblöcke aus der Wissensbasis "
                        f"(Block-Zusammenfassungen):\n\n{summary}\n\n"
                        f"Frage: {question}"
                    ),
                },
            ]
        # leerer Map-Output -> Fallback: messages bleibt die Einzelantwort

    final = chat_with_length_retry(messages)
    violations = validate_answer(final, allowed)
    violations += validate_decision_support_answer(question, final, allowed)
    regenerations = 0
    if violations:
        regenerations = 1
        warn = (
            "Deine letzte Antwort verstieß gegen die Regeln: "
            + "; ".join(violations)
            + f". Erlaubte KB-IDs sind ausschließlich: {', '.join(sorted(set(allowed)))}. "
            "Beantworte die Frage erneut und zitiere nur diese IDs — oder verweigere "
            f"mit dem vorgesehenen Satz („{REFUSAL_MESSAGE}“)."
        )
        retry = chat_with_length_retry(
            messages
            + [{"role": "assistant", "content": final},
               {"role": "user", "content": warn}]
        )
        retry_violations = validate_answer(retry, allowed)
        retry_violations += validate_decision_support_answer(question, retry, allowed)
        if not retry_violations:
            final = retry
            violations = []
        else:
            return finish(
                UNCERTAIN_MESSAGE,
                refused=True,
                verified=False,
                citations=[],
                regenerations=regenerations,
                draft=retry,
                sources=[],
            )

    citations = sorted(set(CITE_RE.findall(final)))
    refused = looks_like_refusal(final)
    sources = [
        {
            "id": cid,
            "title": by_id[cid].title,
            "section": by_id[cid].section,
            "stand": by_id[cid].stand,
            "work": by_id[cid].work,
        }
        for cid in citations
        if cid in by_id
    ]
    return finish(
        final, refused=refused, verified=not violations,
        citations=citations, regenerations=regenerations, sources=sources,
    )