"""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, 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 Werte mit ihren IDs und
   kennzeichne den Widerspruch mit ⚠. 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.

Verletze Regel 2 oder Regel 4 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)-[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 _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):
        return strip_think(
            client.chat(
                cfg.answer_model,
                msgs,
                temperature=cfg.temperature,
                num_ctx=cfg.num_ctx,
                num_predict=cfg.num_predict,
                think=cfg.think,
            )
        )

    if map_messages is not None:
        summary = chat(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(messages)
    violations = validate_answer(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(
            messages
            + [{"role": "assistant", "content": final},
               {"role": "user", "content": warn}]
        )
        retry_violations = validate_answer(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,
    )