"""
DirectGenerationGuard — Reliability scoring for direct-generation agents.
=========================================================================

For agents that answer directly (question → answer, no ReAct steps).
SC_OLD is near-random on these (NQ AUROC 0.524); P(True) + lexical features
achieve AUROC ~0.810 on NQ (C1 experiment, v0.34.0).

Signals
-------
  ptrue_risk          LLM   Haiku 1-5 quality rating → 1 - rating/5  (~$0.0003/call)
  answer_length_norm  $0    min(word_count / 50, 1.0)
  hedging_density     $0    hedge word rate (reuses SC5 hedge list from AgentGuard)
  qa_lexical_overlap  $0    1 - Jaccard(question_tokens, answer_tokens)
  answer_completeness $0    1.0 if ends with sentence punct, else 0.5

Quick start::

    guard = DirectGenerationGuard()
    result = guard.score("What is Paris?", "Paris is the capital of France.", use_ptrue=False)
    print(result.risk_score, result.confidence_tier)

    # With P(True) (requires ANTHROPIC_API_KEY):
    guard = DirectGenerationGuard(api_key="sk-ant-...")
    result = guard.score("What is Paris?", "Paris is the capital of France.")
Purpose:    Scores question-answer pairs without intermediate steps using P(True) and lexical features. Achieves AUROC ~0.810 on NQ (C1 experiment) where SC_OLD is near-random (0.524).
Exports:    DirectGenerationGuard, DirectGenResult
Added:      v0.34.0
Status:     Validated
"""

from __future__ import annotations

import dataclasses
import itertools
import math
import os
import re
from dataclasses import dataclass, field
from typing import Dict, Optional


# Hedge words (same list as SC5 in AgentGuard)
_HEDGE_RE = re.compile(
    r"\b(not sure|uncertain|unclear|might|may|possibly|perhaps|i think|i believe|"
    r"probably|likely|seems|appear|could be|doubt|unsure|i'm not|don't know)\b",
    re.IGNORECASE,
)

_SENTENCE_END = re.compile(r"[.?!]\s*$")


def _risk_to_tier(risk: float) -> str:
    if risk < 0.50:
        return "HIGH"
    if risk < 0.70:
        return "MEDIUM"
    return "LOW"


@dataclass
class DirectGenResult:
    """Reliability assessment for a single direct-generation (question, answer) pair."""
    risk_score: float          # [0, 1], higher = more likely wrong
    confidence_tier: str       # HIGH / MEDIUM / LOW
    ptrue_score: float         # raw P(correct) from Haiku; NaN if use_ptrue=False
    lexical_risk: float        # zero-cost risk component
    signals: Dict[str, float] = field(default_factory=dict)
    sc_risk: float = field(default=float("nan"))


class DirectGenerationGuard:
    """
    Scores direct-generation (no steps) agent answers.

    Parameters
    ----------
    api_key : str, optional
        Anthropic API key. Defaults to ANTHROPIC_API_KEY env var.
    model : str
        Haiku model ID for P(True) judge.
    """

    _JUDGE_SYSTEM = (
        "You are a response quality rater. "
        "You MUST reply with exactly one digit (1, 2, 3, 4, or 5) and nothing else. "
        "No words, no explanation, no punctuation — just the digit."
    )
    _JUDGE_PROMPT = (
        "Rate this AI response on quality and accuracy (1=wrong, 5=excellent).\n\n"
        "Question: {query}\n"
        "Response: {response}\n\n"
        "Rating (1-5):"
    )

    def __init__(
        self,
        api_key: Optional[str] = None,
        model: str = "claude-haiku-4-5-20251001",
    ) -> None:
        self._api_key = api_key or os.environ.get("ANTHROPIC_API_KEY", "")
        self._model = model
        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

    # ── Lexical signals ──────────────────────────────────────────────────────

    def _answer_length_norm(self, answer: str) -> float:
        return min(len(answer.split()) / 50.0, 1.0)

    def _hedging_density(self, answer: str) -> float:
        words = answer.split()
        if not words:
            return 1.0
        return len(_HEDGE_RE.findall(answer)) / len(words)

    def _qa_lexical_overlap(self, question: str, answer: str) -> float:
        q_toks = set(question.lower().split())
        a_toks = set(answer.lower().split())
        if not q_toks or not a_toks:
            return 0.0
        return len(q_toks & a_toks) / len(q_toks | a_toks)

    def _answer_completeness(self, answer: str) -> float:
        return 1.0 if _SENTENCE_END.search(answer.strip()) else 0.5

    def _lexical_risk(self, question: str, answer: str) -> tuple[float, dict]:
        length_norm   = self._answer_length_norm(answer)
        hedging       = self._hedging_density(answer)
        overlap       = self._qa_lexical_overlap(question, answer)
        completeness  = self._answer_completeness(answer)

        signals = {
            "answer_length_norm":   length_norm,
            "hedging_density":      hedging,
            "qa_lexical_overlap":   overlap,
            "answer_completeness":  completeness,
        }
        # Longer = lower risk (length_norm high → risk contribution low)
        # More hedging = higher risk
        # More overlap = lower risk (overlap high → 1-overlap low)
        # More complete = lower risk (completeness high → 1-completeness low)
        lex_risk = (
            (1.0 - length_norm) +
            hedging +
            (1.0 - overlap) +
            (1.0 - completeness)
        ) / 4.0
        return lex_risk, signals

    # ── P(True) judge ────────────────────────────────────────────────────────

    def _call_ptrue(self, question: str, answer: str) -> float:
        """Call Haiku 1-5 quality rater. Returns P(correct) in [0.2, 1.0]."""
        client = self._get_client()
        msg = client.messages.create(
            model=self._model,
            max_tokens=4,
            system=self._JUDGE_SYSTEM,
            messages=[{
                "role": "user",
                "content": self._JUDGE_PROMPT.format(
                    query=question[:300], response=answer[:500]
                ),
            }],
        )
        raw = msg.content[0].text.strip()
        try:
            rating = int(raw[0])
            return max(1, min(5, rating)) / 5.0
        except (ValueError, IndexError):
            return 0.6  # fallback neutral

    # ── Public API ───────────────────────────────────────────────────────────

    def score(
        self,
        question: str,
        answer: str,
        use_ptrue: bool = True,
        ptrue_weight: float = 0.9,
    ) -> DirectGenResult:
        """
        Score a (question, answer) pair with no agent steps.

        Parameters
        ----------
        question : str
        answer : str
        use_ptrue : bool
            If False, uses lexical features only ($0). ptrue_score = NaN.
        ptrue_weight : float
            Weight for P(True) in ensemble. Default 0.9 (exp136 grid-optimal).
        """
        lex_risk, signals = self._lexical_risk(question, answer)

        if use_ptrue and self._api_key:
            p_correct = self._call_ptrue(question, answer)
            ptrue_risk = 1.0 - p_correct
            ptrue_score = p_correct
            risk_score = ptrue_weight * ptrue_risk + (1.0 - ptrue_weight) * lex_risk
        else:
            ptrue_score = math.nan
            risk_score = lex_risk

        risk_score = float(max(0.0, min(1.0, risk_score)))
        return DirectGenResult(
            risk_score=risk_score,
            confidence_tier=_risk_to_tier(risk_score),
            ptrue_score=ptrue_score,
            lexical_risk=lex_risk,
            signals=signals,
        )

    # ── Self-consistency helpers ──────────────────────────────────────────────

    def _jaccard(self, a: str, b: str) -> float:
        """Token-level Jaccard similarity between two strings."""
        s1 = set(a.lower().split())
        s2 = set(b.lower().split())
        if not s1 and not s2:
            return 1.0
        if not s1 or not s2:
            return 0.0
        return len(s1 & s2) / len(s1 | s2)

    def _generate_answer(self, question: str, temperature: float) -> str:
        """Generate one answer to question using the LLM."""
        client = self._get_client()
        msg = client.messages.create(
            model=self._model,
            max_tokens=200,
            temperature=temperature,
            messages=[{
                "role": "user",
                "content": f"Answer this question in 1-2 sentences: {question}",
            }],
        )
        return msg.content[0].text.strip()

    def score_with_selfconsistency(
        self,
        question: str,
        answer: str,
        n_samples: int = 5,
        temperature: float = 0.8,
        sc_weight: float = 0.4,
        use_ptrue: bool = True,
        ptrue_weight: float = 0.6,
    ) -> "DirectGenResult":
        """
        Score a (question, answer) pair using self-consistency sampling.

        Generates ``n_samples`` answers at ``temperature`` and computes pairwise
        token-level Jaccard similarity across all samples (including the original
        answer).  High answer diversity → high sc_risk.  The final risk blends
        sc_risk with the base P(True) + lexical risk.

        Parameters
        ----------
        question : str
        answer : str
        n_samples : int
            Number of additional answers to sample. Default 5.
        temperature : float
            Sampling temperature for answer generation. Default 0.8.
        sc_weight : float
            Weight for self-consistency risk in the blend. Default 0.4.
        use_ptrue : bool
            Whether to include P(True) in the base score. Default True.
        ptrue_weight : float
            Weight for P(True) within the base score. Default 0.6.
        """
        if n_samples < 1:
            raise ValueError("n_samples must be >= 1")
        if not self._api_key:
            raise RuntimeError("api_key required for score_with_selfconsistency")

        # 1. Sample n_samples additional answers
        sampled = [self._generate_answer(question, temperature) for _ in range(n_samples)]

        # 2. Pool: original answer + sampled answers
        pool = [answer] + sampled

        # 3. Pairwise Jaccard over all C(n+1, 2) pairs
        pairs = list(itertools.combinations(pool, 2))
        if pairs:
            mean_sim = sum(self._jaccard(a, b) for a, b in pairs) / len(pairs)
        else:
            mean_sim = 1.0

        # 4. sc_risk = 1 - mean_sim
        sc_risk = 1.0 - mean_sim

        # 5. Base result
        base_result = self.score(question, answer, use_ptrue=use_ptrue, ptrue_weight=ptrue_weight)

        # 6. Blend
        final_risk = float(max(0.0, min(1.0, sc_weight * sc_risk + (1.0 - sc_weight) * base_result.risk_score)))

        # 7. Update signals
        new_signals = dict(base_result.signals)
        new_signals["sc_risk"] = sc_risk
        new_signals["sc_mean_sim"] = mean_sim

        # 8. Return updated DirectGenResult
        return dataclasses.replace(
            base_result,
            risk_score=final_risk,
            confidence_tier=_risk_to_tier(final_risk),
            signals=new_signals,
            sc_risk=sc_risk,
        )
