"""
llm_guard/domain_router.py
==========================

DomainRouter: zero-label, question-type-driven signal selection for
reliable LLM chain scoring across heterogeneous domains.

Validated routing table (exp_E2_domain_router_verify, val split, 2026-03-24):

  Domain Type          Signal           AUROC   95% CI            Confidence
  ---------------------------------------------------------------------------
  fact_verification    faj_r            0.810   [0.681, 0.922]    HIGH
  multi_hop_entity     mfe_sem_sim      0.788   [0.620, 0.924]    HIGH
  trivia_recall        faj_a            0.751   [0.631, 0.858]    HIGH
  open_factual         mfe_sc_pair      0.667   [0.487, 0.833]    LOW (n=40)
  multi_hop_bridge     sc_old           0.625   [0.452, 0.785]    LOW (n=43)

  Mean AUROC: 0.728

Key signals:
  faj_r:        Raw Haiku judge score (~$0.0001/chain). Strong on claim verification.
  mfe_sem_sim:  1 - semantic_sim(Q, A). Zero-cost text feature. Strong on 2Wiki-style
                multi-hop where wrong answers don't address the question.
  faj_a:        Feature-augmented Haiku judge (~$0.0001/chain). Strong on trivia recall.
  mfe_sc_pair:  1 - F1(direct_run0, direct_run1). Zero-cost pairwise consistency.
                Strong on open factual NQ-style questions.
  sc_old:       1 - mean_F1(chain_answer, direct_haiku_answers). Universal fallback.

Usage:
    from llm_guard import DomainRouter, DomainRoutedResult

    router = DomainRouter(anthropic_api_key="sk-ant-...")
    result = router.score(question, steps, final_answer)
    print(result.risk_score, result.domain_type, result.signal_used)
"""
from __future__ import annotations

import os
import re
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional

__all__ = [
    "QuestionType",
    "DomainRoutedResult",
    "QuestionTypeClassifier",
    "DomainRouter",
]

# ── Question type taxonomy ────────────────────────────────────────────────────


class QuestionType(str, Enum):
    """Five question types, each with a validated routing signal."""
    FACT_VERIFICATION  = "fact_verification"   # claim truth check → faj_r
    MULTI_HOP_ENTITY   = "multi_hop_entity"    # entity bridging/comparison → mfe_sem_sim
    TRIVIA_RECALL      = "trivia_recall"        # obscure recall → faj_a
    OPEN_FACTUAL       = "open_factual"         # single-hop factual → mfe_sc_pair
    MULTI_HOP_BRIDGE   = "multi_hop_bridge"     # predicate-level bridging → sc_old


# ── Result dataclass ──────────────────────────────────────────────────────────


@dataclass
class DomainRoutedResult:
    """
    Result from DomainRouter.score().

    Attributes
    ----------
    risk_score : float
        Risk in [0, 1]. Higher = more likely incorrect.
    domain_type : QuestionType
        Detected question type.
    signal_used : str
        Name of the signal that produced risk_score.
    signal_auroc : float
        Validated full-split AUROC for this signal on this domain type.
    confidence : str
        "HIGH" | "MEDIUM" | "LOW" — statistical confidence tier from E2 verification.
    raw_signal_value : float
        The raw (pre-normalised) signal value, for debugging.
    classifier_source : str
        "regex" or "haiku" — how the domain type was determined.
    extras : dict
        Additional per-signal metadata (e.g., judge response text).
    """
    risk_score: float
    domain_type: QuestionType
    signal_used: str
    signal_auroc: float
    confidence: str
    raw_signal_value: float
    classifier_source: str
    extras: Dict = field(default_factory=dict)


# ── Validated routing table ───────────────────────────────────────────────────

_ROUTING_TABLE: Dict[QuestionType, Dict] = {
    QuestionType.FACT_VERIFICATION: {
        "signal": "faj_r",
        "auroc":  0.810,
        "ci":     (0.681, 0.922),
        "confidence": "HIGH",
        "cost":   "~$0.0001/chain",
    },
    QuestionType.MULTI_HOP_ENTITY: {
        "signal": "mfe_sem_sim",
        "auroc":  0.788,
        "ci":     (0.620, 0.924),
        "confidence": "HIGH",
        "cost":   "$0",
    },
    QuestionType.TRIVIA_RECALL: {
        "signal": "faj_a",
        "auroc":  0.751,
        "ci":     (0.631, 0.858),
        "confidence": "HIGH",
        "cost":   "~$0.0001/chain",
    },
    QuestionType.OPEN_FACTUAL: {
        "signal": "mfe_sc_pair",
        "auroc":  0.667,
        "ci":     (0.487, 0.833),
        "confidence": "LOW",
        "cost":   "$0 (cached)",
    },
    QuestionType.MULTI_HOP_BRIDGE: {
        "signal": "sc_old",
        "auroc":  0.625,
        "ci":     (0.452, 0.785),
        "confidence": "LOW",
        "cost":   "~$0.0002/chain",
    },
}


# ── Regex classifier ──────────────────────────────────────────────────────────

# Fact-verification: claim-like phrasing, true/false framing
_FEVER_PATTERNS = re.compile(
    r"\b(true|false|correct|incorrect|fact|claim|verify|stated|assert|says that|"
    r"it is (true|false)|is it true|is this (true|correct))\b",
    re.IGNORECASE,
)
# Multi-hop entity: comparison, bridging entity questions
_ENTITY_HOP_PATTERNS = re.compile(
    r"\b(both|same|share|compare|common|as well as|in common|both .{1,30} and|"
    r"which .{1,50} also|who .{1,50} also|what .{1,60} as well)\b",
    re.IGNORECASE,
)
# Multi-hop bridge: chained entity resolution "who directed the film that starred..."
_BRIDGE_HOP_PATTERNS = re.compile(
    r"\b(who (directed|wrote|produced|starred|played|founded|created|built|"
    r"invented|designed|led|commanded|married|killed) (the|a|an) .{1,60} "
    r"(that|which|who|starring|featuring|based on))\b"
    r"|(\bwhich .{1,50} (that|which) .{1,50} (of|in|at|for|from)\b)",
    re.IGNORECASE,
)
# Trivia: specific dates, numbers, proper nouns with obscure detail
_TRIVIA_PATTERNS = re.compile(
    r"\b(in what (year|month|century)|on what date|how many .{1,30} (did|does|has|have|"
    r"were|was|is|are)|what (was|is) the (first|last|only|original|main|largest|smallest|"
    r"oldest|youngest)|who (won|lost|received|became|served as|was known as))\b",
    re.IGNORECASE,
)


def _classify_by_regex(question: str) -> tuple[QuestionType, float]:
    """
    Classify question type by regex patterns.
    Returns (QuestionType, confidence) where confidence is 0.0-1.0.
    Low-confidence (< 0.6) should trigger Haiku fallback if available.
    """
    q = question.strip()

    if _FEVER_PATTERNS.search(q):
        return QuestionType.FACT_VERIFICATION, 0.85

    # Multi-hop bridge has priority over entity hop (more specific pattern)
    if _BRIDGE_HOP_PATTERNS.search(q):
        return QuestionType.MULTI_HOP_BRIDGE, 0.75

    if _ENTITY_HOP_PATTERNS.search(q):
        return QuestionType.MULTI_HOP_ENTITY, 0.75

    if _TRIVIA_PATTERNS.search(q):
        return QuestionType.TRIVIA_RECALL, 0.70

    # Default: open_factual for single-hop wh-questions
    return QuestionType.OPEN_FACTUAL, 0.55


# ── QuestionTypeClassifier ────────────────────────────────────────────────────

_CLASSIFY_SYSTEM = (
    "You are a question type classifier. Classify the question into EXACTLY ONE of:\n"
    "  fact_verification  — a claim to be verified as true/false\n"
    "  multi_hop_entity   — requires bridging/comparing two entities\n"
    "  trivia_recall      — requires obscure factual recall\n"
    "  open_factual       — single-hop factual lookup\n"
    "  multi_hop_bridge   — requires chaining entity properties\n\n"
    "Respond with ONLY the type name, nothing else."
)

_TYPE_MAP = {
    "fact_verification": QuestionType.FACT_VERIFICATION,
    "multi_hop_entity":  QuestionType.MULTI_HOP_ENTITY,
    "trivia_recall":     QuestionType.TRIVIA_RECALL,
    "open_factual":      QuestionType.OPEN_FACTUAL,
    "multi_hop_bridge":  QuestionType.MULTI_HOP_BRIDGE,
}


class QuestionTypeClassifier:
    """
    Classifies questions into QuestionType using regex (fast, free) with optional
    Haiku fallback for low-confidence cases.

    Parameters
    ----------
    use_llm_fallback : bool
        If True, calls Haiku when regex confidence < 0.65. Default False.
    api_key : str, optional
        Anthropic API key. If not provided, reads ANTHROPIC_API_KEY env var.
    """

    def __init__(
        self,
        use_llm_fallback: bool = False,
        api_key: Optional[str] = None,
    ):
        self.use_llm_fallback = use_llm_fallback
        self._api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        self._client = None

    def _get_client(self):
        if self._client is None:
            import anthropic
            self._client = anthropic.Anthropic(api_key=self._api_key)
        return self._client

    def _classify_haiku(self, question: str) -> Optional[QuestionType]:
        """Call Haiku to classify. Returns None on failure."""
        try:
            resp = self._get_client().messages.create(
                model="claude-haiku-4-5-20251001",
                max_tokens=20,
                system=_CLASSIFY_SYSTEM,
                messages=[{"role": "user", "content": question[:400]}],
            )
            text = resp.content[0].text.strip().lower()
            return _TYPE_MAP.get(text)
        except Exception:
            time.sleep(0.5)
            return None

    def classify(self, question: str) -> tuple[QuestionType, str]:
        """
        Classify question type.

        Returns
        -------
        (QuestionType, source) where source is "regex" or "haiku".
        """
        qtype, conf = _classify_by_regex(question)

        if self.use_llm_fallback and conf < 0.65 and self._api_key:
            haiku_type = self._classify_haiku(question)
            if haiku_type is not None:
                return haiku_type, "haiku"

        return qtype, "regex"


# ── Signal computers ──────────────────────────────────────────────────────────

_STOP_WORDS = {
    "the", "a", "an", "is", "are", "was", "were", "be", "been", "being",
    "have", "has", "had", "do", "does", "did", "will", "would", "could",
    "should", "may", "might", "shall", "can", "this", "that", "these", "those",
    "it", "its", "they", "them", "their", "of", "in", "on", "at", "to", "for",
    "with", "by", "from", "up", "out", "as", "into", "through", "and", "or",
    "but", "not", "so", "if", "then", "than", "about", "which", "who", "what",
    "how", "when", "where",
}

_JUDGE_SYSTEM = (
    "You are an answer quality judge. Given a question and an answer, rate the "
    "answer quality on a scale from 1 to 5 where:\n"
    "  1 = completely wrong or irrelevant\n"
    "  2 = mostly wrong with minor correct elements\n"
    "  3 = partially correct\n"
    "  4 = mostly correct with minor issues\n"
    "  5 = completely correct and well-stated\n\n"
    "Respond with ONLY the number."
)

_DIRECT_SYSTEM = (
    "You are a knowledgeable assistant. Answer the question directly and concisely "
    "in 1-2 sentences using your knowledge. Do NOT say 'I don't know' — give your "
    "best answer. Only output the answer, no preamble."
)


def _tokenize(text: str) -> set:
    return {w for w in re.findall(r"[a-zA-Z]+", text.lower()) if w not in _STOP_WORDS}


def _token_f1(pred: str, ref: str) -> float:
    p, r = _tokenize(pred), _tokenize(ref)
    if not p or not r:
        return 0.0
    c = p & r
    prec, rec = len(c) / len(p), len(c) / len(r)
    return 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0


def _chain_context(steps: List[Dict], max_chars: int = 800) -> str:
    """Extract observation text from ReAct steps."""
    parts = []
    total = 0
    for step in steps:
        obs = str(step.get("observation", "") or "").strip()
        if len(obs) < 10:
            continue
        snippet = obs[:300]
        if total + len(snippet) > max_chars:
            break
        parts.append(snippet)
        total += len(snippet)
    return " ".join(parts)


def _compute_mfe_sem_sim(question: str, final_answer: str, context: str) -> float:
    """
    MFE Phase A f1 (sem_sim_qa): semantic similarity between question and answer.
    Returns risk = 1 - sim (higher = less grounded in question).
    Zero cost — computed from text only.
    """
    try:
        from llm_guard.features.mechanism_layer import MechanismFeatureExtractor
        _mfe = MechanismFeatureExtractor()
        feats = _mfe.extract_features(question, final_answer, context)
        sem_sim = float(feats[0])  # f1 = sem_sim_qa
        return 1.0 - sem_sim
    except Exception:
        return 0.5  # neutral fallback


def _compute_mfe_sc_pair(
    question: str, final_answer: str, api_key: str, n_runs: int = 2
) -> tuple[float, dict]:
    """
    Phase B proxy: 1 - F1(direct_run0, direct_run1).
    Higher = more inconsistent direct answers = higher uncertainty.
    Requires API calls if not cached.
    """
    try:
        import anthropic
        client = anthropic.Anthropic(api_key=api_key)
        answers = []
        for _ in range(n_runs):
            resp = client.messages.create(
                model="claude-haiku-4-5-20251001",
                max_tokens=80,
                system=_DIRECT_SYSTEM,
                messages=[{"role": "user", "content": question[:400]}],
            )
            answers.append(resp.content[0].text.strip())
            time.sleep(0.1)

        if len(answers) == 2:
            f1 = _token_f1(answers[0], answers[1])
            return 1.0 - f1, {"direct_answers": answers}
        return 0.5, {}
    except Exception:
        return 0.5, {}


def _compute_faj(
    question: str, final_answer: str, context: str, api_key: str, augmented: bool
) -> tuple[float, dict]:
    """
    FAJ-R (raw Haiku judge) or FAJ-A (feature-augmented judge).
    Returns (risk, extras).
    """
    try:
        import anthropic
        client = anthropic.Anthropic(api_key=api_key)

        if augmented:
            # FAJ-A: provide context to judge
            prompt = (f"Question: {question[:300]}\n\n"
                      f"Evidence: {context[:400]}\n\n"
                      f"Answer: {final_answer[:200]}")
        else:
            # FAJ-R: no context
            prompt = f"Question: {question[:300]}\n\nAnswer: {final_answer[:200]}"

        resp = client.messages.create(
            model="claude-haiku-4-5-20251001",
            max_tokens=5,
            system=_JUDGE_SYSTEM,
            messages=[{"role": "user", "content": prompt}],
        )
        text = resp.content[0].text.strip()
        match = re.search(r"[1-5]", text)
        if match:
            score = int(match.group()) / 5.0  # 1→0.2, 5→1.0
            risk = 1.0 - score               # 1→0.8, 5→0.0
            return risk, {"judge_score": int(match.group()), "judge_text": text}
        return 0.5, {"judge_text": text}
    except Exception:
        return 0.5, {}


def _compute_sc_old(
    question: str, final_answer: str, api_key: str, n_runs: int = 2
) -> tuple[float, dict]:
    """SC_OLD: 1 - mean_F1(final_answer, direct_haiku_answers)."""
    try:
        import anthropic
        client = anthropic.Anthropic(api_key=api_key)
        f1s = []
        direct_answers = []
        for _ in range(n_runs):
            resp = client.messages.create(
                model="claude-haiku-4-5-20251001",
                max_tokens=80,
                system=_DIRECT_SYSTEM,
                messages=[{"role": "user", "content": question[:400]}],
            )
            direct = resp.content[0].text.strip()
            direct_answers.append(direct)
            f1s.append(_token_f1(final_answer, direct))
            time.sleep(0.1)

        risk = float(1.0 - (sum(f1s) / len(f1s))) if f1s else 0.5
        return risk, {"direct_answers": direct_answers, "f1_scores": f1s}
    except Exception:
        return 0.5, {}


# ── DomainRouter ──────────────────────────────────────────────────────────────


class DomainRouter:
    """
    Routes agent chains to the optimal scoring signal based on detected question type.

    Validated routing table (exp_E2_domain_router_verify, 2026-03-24):
      fact_verification → faj_r      (AUROC 0.810, HIGH confidence)
      multi_hop_entity  → mfe_sem_sim (AUROC 0.788, HIGH confidence)
      trivia_recall     → faj_a      (AUROC 0.751, HIGH confidence)
      open_factual      → mfe_sc_pair (AUROC 0.667, LOW confidence — n=40)
      multi_hop_bridge  → sc_old     (AUROC 0.625, LOW confidence — n=43)

    Parameters
    ----------
    api_key : str, optional
        Anthropic API key for signals that require LLM calls (faj_r, faj_a, sc_old,
        mfe_sc_pair). If not provided, reads ANTHROPIC_API_KEY env var.
        mfe_sem_sim is zero-cost and does not require an API key.
    use_llm_classifier : bool
        If True, uses Haiku for domain classification when regex confidence < 0.65.
        Default False (regex only, faster and free).

    Examples
    --------
    >>> router = DomainRouter(api_key="sk-ant-...")
    >>> result = router.score(
    ...     question="Which band did the guitarist of Radiohead join?",
    ...     steps=[{"action": "Search", "observation": "Thom Yorke founded..."}],
    ...     final_answer="Thom Yorke"
    ... )
    >>> print(result.risk_score, result.domain_type, result.signal_used)
    0.72 multi_hop_bridge sc_old
    """

    ROUTING_TABLE = _ROUTING_TABLE

    def __init__(
        self,
        api_key: Optional[str] = None,
        use_llm_classifier: bool = False,
    ):
        self._api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        self.classifier = QuestionTypeClassifier(
            use_llm_fallback=use_llm_classifier,
            api_key=self._api_key,
        )

    def routing_table(self) -> Dict[str, Dict]:
        """Return the validated routing table as a plain dict."""
        return {
            qtype.value: {
                "signal": info["signal"],
                "auroc": info["auroc"],
                "ci": list(info["ci"]),
                "confidence": info["confidence"],
                "cost": info["cost"],
            }
            for qtype, info in _ROUTING_TABLE.items()
        }

    def score(
        self,
        question: str,
        steps: List[Dict],
        final_answer: str,
        domain_type: Optional[QuestionType] = None,
    ) -> DomainRoutedResult:
        """
        Score a completed agent chain using the domain-appropriate signal.

        Parameters
        ----------
        question : str
        steps : list of dicts
            ReAct-format steps with 'action' and 'observation' keys.
        final_answer : str
        domain_type : QuestionType, optional
            Override automatic classification. If None, uses QuestionTypeClassifier.

        Returns
        -------
        DomainRoutedResult
        """
        # Step 1: classify domain
        if domain_type is not None:
            qtype = domain_type
            source = "override"
        else:
            qtype, source = self.classifier.classify(question)

        route = _ROUTING_TABLE[qtype]
        signal_name = route["signal"]
        context = _chain_context(steps)

        # Step 2: compute signal
        if signal_name == "mfe_sem_sim":
            raw, extras = _compute_mfe_sem_sim(question, final_answer, context), {}
        elif signal_name == "mfe_sc_pair":
            raw, extras = _compute_mfe_sc_pair(question, final_answer, self._api_key)
        elif signal_name == "faj_r":
            raw, extras = _compute_faj(question, final_answer, context,
                                       self._api_key, augmented=False)
        elif signal_name == "faj_a":
            raw, extras = _compute_faj(question, final_answer, context,
                                       self._api_key, augmented=True)
        elif signal_name == "sc_old":
            raw, extras = _compute_sc_old(question, final_answer, self._api_key)
        else:
            raw, extras = 0.5, {}

        return DomainRoutedResult(
            risk_score=round(float(raw), 4),
            domain_type=qtype,
            signal_used=signal_name,
            signal_auroc=route["auroc"],
            confidence=route["confidence"],
            raw_signal_value=round(float(raw), 4),
            classifier_source=source,
            extras=extras,
        )
