"""Output guardrails.

The fine-tune teaches the model to behave; guardrails are the deterministic safety net that runs
*after* generation:
  - require an inline citation when we actually retrieved context,
  - block personalized financial/"which plan should I buy" advice,
  - escalate to a human when we have no grounded answer.
"""
from __future__ import annotations

import re
from dataclasses import dataclass

from .prompt import RetrievedChunk
from .settings import Config

# Phrases that signal the model is giving a personalized recommendation rather than facts.
_ADVICE_PATTERNS = [
    r"\byou should (?:buy|choose|purchase|get|pick|go with)\b",
    r"\bi recommend (?:that you|you)\b",
    r"\bthe best (?:plan|policy|option) for you\b",
    r"\bi'd (?:suggest|advise) you\b",
]
# A valid citation is a numbered passage marker like [1]; a bare "[All]"/"[source]" doesn't count,
# so a model's "I can't answer from the context" refusal is correctly treated as an escalation.
_CITATION_RE = re.compile(r"\[\d+\]")


@dataclass
class GuardrailResult:
    answer: str
    ok: bool
    reason: str = ""


class Guardrails:
    def __init__(self, cfg: Config):
        self.require_citation = bool(cfg.get("guardrails.require_citation", True))
        self.escalation = cfg.get(
            "guardrails.escalation_message",
            "I can't answer that confidently. Please contact a licensed advisor.",
        ).strip()
        self._advice_res = [re.compile(p, re.IGNORECASE) for p in _ADVICE_PATTERNS]

    def check(self, answer: str, chunks: list[RetrievedChunk]) -> GuardrailResult:
        text = (answer or "").strip()

        # No grounded context -> escalate rather than let the model free-associate.
        if not chunks:
            return GuardrailResult(self.escalation, ok=False, reason="no_context")

        if not text:
            return GuardrailResult(self.escalation, ok=False, reason="empty_answer")

        # Personalized advice -> replace with a factual escalation.
        for rx in self._advice_res:
            if rx.search(text):
                return GuardrailResult(
                    self.escalation
                    + " (I can explain the differences between plans, but I can't tell you which "
                    "one to choose.)",
                    ok=False,
                    reason="personalized_advice",
                )

        # Must cite a source when context was available.
        if self.require_citation and not _CITATION_RE.search(text):
            return GuardrailResult(self.escalation, ok=False, reason="missing_citation")

        return GuardrailResult(text, ok=True)