"""
MultiTurnGuard — Reliability scoring for multi-turn conversation agents.
========================================================================

Scores conversations using a sliding window over assistant turns.
Stateless: caller passes the full conversation each call.
Composes with GuardState for persistence.

Validation status (v0.37.0)
---------------------------
  use_ptrue=True : AUROC 0.757 [0.628, 0.875] on MT-Bench human judgments, n=50.
                   CI is wide (±0.124). Full 3.3k-pair evaluation pending.
  use_ptrue=False: Lexical-only mode. NOT validated for real-world multi-turn
                   quality discrimination. Returns scores but they carry no
                   empirical reliability guarantee. A UserWarning is emitted.

Quick start (validated mode)::

    from llm_guard import MultiTurnGuard

    guard = MultiTurnGuard(api_key="sk-ant-...")  # or set ANTHROPIC_API_KEY
    conv = [
        {"role": "user", "content": "What is the capital of France?"},
        {"role": "assistant", "content": "Paris is the capital of France."},
    ]
    result = guard.score_turn(conv, use_ptrue=True)
    print(result.risk_score, result.drift_detected)
Purpose:    Scores conversations using a sliding window over assistant turns with P(True) or lexical features. Validated AUROC 0.757 [0.628, 0.875] on MT-Bench human judgments with use_ptrue=True.
Exports:    MultiTurnGuard, MultiTurnResult
Added:      v0.1.0
Status:     Validated
"""

from __future__ import annotations

import os
import re
import warnings
from dataclasses import dataclass
from typing import List, Optional


_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 _lexical_quality(question: str, answer: str) -> float:
    """Return quality score in [0, 1] using zero-cost lexical signals."""
    words = answer.split()
    n = max(len(words), 1)

    length_score   = min(len(words) / 20.0, 1.0)                      # longer = better
    hedge_penalty  = len(_HEDGE_RE.findall(answer)) / n                # more hedging = worse
    complete_bonus = 0.2 if _SENTENCE_END.search(answer.strip()) else 0.0

    quality = max(0.0, min(1.0,
        0.4 * length_score +
        0.4 * (1.0 - hedge_penalty) +
        complete_bonus  # 0.0 or 0.2 sentence-ending bonus
    ))
    return quality


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


@dataclass
class MultiTurnResult:
    """Reliability assessment for a multi-turn conversation window."""
    risk_score: float        # 1.0 - mean_quality
    confidence_tier: str     # HIGH / MEDIUM / LOW
    drift_detected: bool     # True if quality slope < -0.10 over window >= 3
    turn_scores: List[float] # per-turn quality scores (len = window_size)
    window_size: int         # actual number of turns scored
    quality_preference: str = ""
    """Qualitative label: 'excellent'(>=0.8), 'good'(>=0.6), 'degraded'(>=0.4), 'poor'(<0.4).
    Empty string for empty conversations. Measures MT-Bench preference quality, NOT factual correctness."""
    is_factual_reliability_signal: bool = False
    """Always False. quality_preference measures MT-Bench response-quality preference,
    NOT factual correctness or chain-level reliability. Do NOT use this field or
    quality_preference to make deployment decisions about factual accuracy.
    For factual reliability use AgentGuard.score_chain() or score_with_ptrue()."""
    construct: str = ""
    """Which reliability construct this result measures.
      "quality_preference"    — score_turn(): MT-Bench response-quality preference
                                (AUROC 0.667, MT-Bench n=2575)
      "factual_reliability"   — score_factual_turn(): P(True) factual correctness
                                (uses same P(True) prompt as AgentGuard)
      ""                      — internal/unset (treat as quality_preference)
    """


class MultiTurnGuard:
    """
    Scores multi-turn conversations using a sliding window of assistant turns.

    Stateless: no internal state. Caller passes full conversation each call.

    Parameters
    ----------
    api_key : str, optional
        Anthropic API key for P(True) judge.
    model : str
        Haiku model for P(True) quality rating.

    Validated Performance
    --------------------
    AUROC 0.667 [0.652, 0.681] on lmsys/mt_bench_human_judgments (n=2,575 preference
    pairs, exp_multiturn_fullbench, 2026-03-20).  Use ``use_ptrue=True`` for the
    validated signal path.

    Construct Validity Note
    -----------------------
    MultiTurnGuard measures **response quality preference** as judged by MT-Bench
    human raters — not factual correctness.  An assistant response can be fluent,
    well-structured, and preferred by raters while still being factually wrong on
    specific claims.  Do NOT treat MultiTurnGuard scores as equivalent to chain-level
    factual reliability (AgentGuard).  Use for:

    - Detecting degraded responses (coherence failures, refusals, off-topic replies)
    - Alerting on consistent quality drop within a session (drift_detected flag)
    - Pre-filtering low-quality completions before downstream processing

    Do NOT use for:
    - Factual correctness verification
    - Cross-domain reliability benchmarking
    - Replacing execution-based code correctness checks

    Zero-Shot Mode (Locked)
    -----------------------
    Few-shot prompting was tested in exp_multiturn_fewshot (2026-03-20) with 3
    in-context examples.  AUROC dropped from 0.669 to 0.607 (delta=-0.062) on
    n=200 MT-Bench pairs.  Zero-shot mode is therefore locked as the only
    supported configuration.  Do NOT attempt to pass few-shot examples; the
    judge prompt is fixed and cannot be overridden.

    Result Fields
    -------------
    ``MultiTurnResult.quality_preference`` — qualitative label derived from
    mean turn quality: ``"excellent"`` (>=0.8), ``"good"`` (>=0.6),
    ``"degraded"`` (>=0.4), ``"poor"`` (<0.4).  Empty string for empty
    conversations.  Measures MT-Bench preference quality, NOT factual
    correctness.
    """

    # Zero-shot rating prompts — locked after exp_multiturn_fewshot showed
    # few-shot hurts: delta=-0.062 on MT-Bench n=200 pairs (2026-03-20).
    _JUDGE_SYSTEM = (
        "You are a response quality rater. "
        "You MUST reply with exactly one digit (1, 2, 3, 4, or 5) and nothing else."
    )
    _JUDGE_PROMPT = (
        "Rate this AI response on quality and accuracy (1=wrong, 5=excellent).\n\n"
        "Question: {query}\nResponse: {response}\n\nRating (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
        if not self._api_key:
            warnings.warn(
                "MultiTurnGuard: no Anthropic API key found. "
                "Calls to score_turn() will fall back to unvalidated lexical-only mode. "
                "Set api_key= or ANTHROPIC_API_KEY and use use_ptrue=True for validated scoring.",
                UserWarning,
                stacklevel=2,
            )

    @classmethod
    def zero_shot(
        cls,
        api_key: Optional[str] = None,
        model: str = "claude-haiku-4-5-20251001",
    ) -> "MultiTurnGuard":
        """
        Named constructor — explicitly signals zero-shot-only mode.

        Identical to ``MultiTurnGuard(api_key=..., model=...)``.  Provided as a
        self-documenting alternative that makes the zero-shot constraint visible
        at the call site and prevents accidental future addition of few-shot
        examples (which hurt: delta=-0.062 on MT-Bench n=200, exp_multiturn_fewshot).

        Example
        -------
        guard = MultiTurnGuard.zero_shot(api_key="sk-ant-...")
        result = guard.score_turn(history, use_ptrue=True)
        """
        return cls(api_key=api_key, model=model)

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

    def _ptrue_quality(self, question: str, answer: str) -> float:
        """Call Haiku 1-5 quality rater. Returns quality in [0.2, 1.0]."""
        try:
            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[:200], response=answer[:400]
                )}],
            )
            raw = msg.content[0].text.strip()
            return max(1, min(5, int(raw[0]))) / 5.0
        except Exception:
            return 0.6  # neutral fallback

    def _extract_pairs(
        self, conversation: List[dict], window: int
    ) -> List[tuple]:
        """
        Extract last `window` (user_turn, assistant_turn) pairs.
        If no user precedes an assistant turn, uses empty string as question.
        """
        pairs = []
        prev_user = ""
        for msg in conversation:
            role = msg.get("role", "")
            content = str(msg.get("content", ""))
            if role == "user":
                prev_user = content
            elif role == "assistant":
                pairs.append((prev_user, content))
                prev_user = ""
        return pairs[-window:] if window > 0 else pairs

    def _score_from_turn_scores(
        self, turn_scores: List[float], window: int = 0
    ) -> MultiTurnResult:
        """Build a MultiTurnResult from pre-computed quality scores."""
        if not turn_scores:
            return MultiTurnResult(
                risk_score=0.5,
                confidence_tier="MEDIUM",
                drift_detected=False,
                turn_scores=[],
                window_size=0,
                quality_preference="",
            )

        mean_q = sum(turn_scores) / len(turn_scores)
        risk = 1.0 - mean_q

        # Drift: slope of quality over window
        drift = False
        if len(turn_scores) >= 3:
            try:
                from scipy.stats import linregress
                xs = list(range(len(turn_scores)))
                slope, *_ = linregress(xs, turn_scores)
                drift = bool(slope < -0.10)
            except ImportError:
                # Fallback: simple first-vs-last comparison
                drift = (turn_scores[-1] - turn_scores[0]) < -0.20

        # Tier overrides for very low/high quality
        if mean_q < 0.4:
            tier = "LOW"
        elif mean_q >= 0.6:
            tier = "HIGH"
        else:
            tier = _risk_to_tier(risk)

        # Quality preference label
        if mean_q >= 0.8:
            quality_preference = "excellent"
        elif mean_q >= 0.6:
            quality_preference = "good"
        elif mean_q >= 0.4:
            quality_preference = "degraded"
        else:
            quality_preference = "poor"

        return MultiTurnResult(
            risk_score=float(max(0.0, min(1.0, risk))),
            confidence_tier=tier,
            drift_detected=drift,
            turn_scores=list(turn_scores),
            window_size=len(turn_scores),
            quality_preference=quality_preference,
        )

    def score_turn(
        self,
        conversation: List[dict],
        window: int = 5,
        use_ptrue: bool = True,
    ) -> MultiTurnResult:
        """
        Score the current conversation state.

        Parameters
        ----------
        conversation : list of dicts
            [{"role": "user"|"assistant", "content": "..."}]
        window : int
            Number of most recent assistant turns to score.
        use_ptrue : bool
            If True, calls Haiku quality rater per turn (requires api_key).
            If False, uses zero-cost lexical signals only. NOT validated for
            real-world multi-turn quality discrimination — a UserWarning is emitted.
        """
        # Surface MT-Bench caveat at the API level (not just in README):
        # quality_preference is a response-quality preference signal from MT-Bench style
        # evaluation, NOT a factual correctness or chain-level reliability signal.
        warnings.warn(
            "MultiTurnGuard.score_turn(): the returned quality_preference field measures "
            "MT-Bench response-quality preference (AUROC=0.667 on MT-Bench n=2575). "
            "It does NOT measure factual correctness. Do not use quality_preference "
            "for deployment decisions about factual accuracy. "
            "For factual reliability use AgentGuard.score_chain() or score_with_ptrue().",
            UserWarning,
            stacklevel=2,
        )
        if not use_ptrue:
            warnings.warn(
                "MultiTurnGuard.score_turn() called with use_ptrue=False. "
                "Lexical-only mode is not validated for real-world multi-turn quality "
                "discrimination (no empirical AUROC on external benchmarks). "
                "Pass use_ptrue=True with a valid Anthropic API key for validated scoring. "
                "Do not publish results from lexical-only mode as reliability estimates.",
                UserWarning,
                stacklevel=2,
            )
        elif not self._api_key:
            warnings.warn(
                "MultiTurnGuard.score_turn(): use_ptrue=True requested but no API key is set. "
                "Falling back to unvalidated lexical-only mode. "
                "Set ANTHROPIC_API_KEY or pass api_key= to MultiTurnGuard().",
                UserWarning,
                stacklevel=2,
            )

        pairs = self._extract_pairs(conversation, window)
        if not pairs:
            return self._score_from_turn_scores([], window)

        turn_scores = []
        for question, answer in pairs:
            if use_ptrue and self._api_key:
                q = self._ptrue_quality(question, answer)
            else:
                q = _lexical_quality(question, answer)
            turn_scores.append(q)

        result = self._score_from_turn_scores(turn_scores, window)
        from dataclasses import replace as _dc_replace
        return _dc_replace(result, construct="quality_preference")

    def _ptrue_factual(self, question: str, answer: str) -> float:
        """
        Call Haiku with the P(True) factual-correctness prompt (1-5 rating).
        Different from _ptrue_quality which uses MT-Bench response-quality prompt.
        Returns correctness score in [0.0, 1.0] — higher = more likely correct.
        """
        _FACTUAL_SYSTEM = "You are a helpful assistant that rates answer correctness."
        _FACTUAL_PROMPT = (
            "Question: {question}\n\n"
            "Answer: {answer}\n\n"
            "Rate how likely this answer is factually correct on a scale of 1 to 5:\n"
            "1=definitely wrong, 2=probably wrong, 3=uncertain, "
            "4=probably correct, 5=definitely correct.\n\n"
            "Reply with only the single digit (1, 2, 3, 4, or 5)."
        )
        try:
            import re as _re
            client = self._get_client()
            msg = client.messages.create(
                model=self._model,
                max_tokens=4,
                system=_FACTUAL_SYSTEM,
                messages=[{"role": "user", "content": _FACTUAL_PROMPT.format(
                    question=question[:300], answer=answer[:400]
                )}],
            )
            raw = msg.content[0].text.strip()
            m = _re.search(r"[1-5]", raw)
            if m:
                return (int(m.group()) - 1) / 4.0  # 1→0.0, 5→1.0
            return 0.5
        except Exception:
            return 0.5

    def score_factual_turn(
        self,
        conversation: List[dict],
        window: int = 5,
    ) -> "MultiTurnResult":
        """
        Score the factual correctness of recent assistant turns using P(True).

        Unlike score_turn() which measures MT-Bench response-quality preference,
        this method asks: "Is each answer factually correct?" using the same
        P(True) prompt validated in AgentGuard (AUROC 0.767 on NQ, 0.810 on NQ,
        exp_ptrue_expanded / exp_nq_ptrue).

        Requires an Anthropic API key. Without an API key a UserWarning is emitted
        and lexical signals are used (not validated for factual accuracy).

        Parameters
        ----------
        conversation : list of dicts
            [{"role": "user"|"assistant", "content": "..."}]
        window : int
            Number of most recent assistant turns to score.

        Returns
        -------
        MultiTurnResult
            Same type as score_turn() for API compatibility.
            quality_preference label reflects factual correctness estimate,
            NOT MT-Bench preference quality.
            is_factual_reliability_signal is True for this method only.
        """
        if not self._api_key:
            warnings.warn(
                "MultiTurnGuard.score_factual_turn(): no API key — falling back to "
                "unvalidated lexical signals. Factual correctness requires an Anthropic "
                "API key. Set ANTHROPIC_API_KEY or pass api_key= to MultiTurnGuard().",
                UserWarning,
                stacklevel=2,
            )

        pairs = self._extract_pairs(conversation, window)
        if not pairs:
            return self._score_from_turn_scores([], window)

        turn_scores = []
        for question, answer in pairs:
            if self._api_key:
                score = self._ptrue_factual(question, answer)
            else:
                score = _lexical_quality(question, answer)
            turn_scores.append(score)

        result = self._score_from_turn_scores(turn_scores, window)
        # Override: this path IS a factual reliability signal, not MT-Bench quality
        from dataclasses import replace as _dc_replace
        return _dc_replace(
            result,
            is_factual_reliability_signal=True,
            construct="factual_reliability",
        )
