"""
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.

Quick start::

    from llm_guard import MultiTurnGuard

    guard = MultiTurnGuard()
    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=False)
    print(result.risk_score, result.drift_detected)
"""

from __future__ import annotations

import os
import re
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


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.
    """

    _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

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

        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)

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

    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.
        """
        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)

        return self._score_from_turn_scores(turn_scores, window)
