"""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 .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. Beantworten die Kontextblöcke die Frage nicht, antworte 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).

Verletze Regel 2 oder Regel 4 niemals — im Zweifel verweigere die Antwort."""

CITE_RE = re.compile(r"\b(?:lb|wk)-[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 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: Retrieval -> Prompt -> LLM -> Post-Validierung."""
    t0 = time.perf_counter()
    own_retriever = retriever is None
    if retriever is None:
        retriever = Retriever(cfg)
    try:
        results = retriever.search(question, n_entries=top_k)
    finally:
        if own_retriever:
            retriever.close()

    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,
        }

    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}

    if client is None:
        client = OllamaClient(
            cfg.ollama_url,
            embed_timeout_s=cfg.embed_timeout_s,
            chat_timeout_s=cfg.chat_timeout_s,
        )
    messages = [
        {"role": "system", "content": 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,
            )
        )

    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,
    )