"""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. 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; stattdätzlich gilt …“ — mit Beleg [ID].
   Verweigere in diesem Fall nicht.

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

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: 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()
    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,
        }

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