"""
llm_guard/intra_chain_drift_detector.py — Intra-Chain Drift Detector
=====================================================================
Detects distributional drift within a single agent chain by wrapping
the existing DriftDetector with a per-chain baseline initialization.

The key insight: long-horizon agents (15-200 steps) experience intra-chain
distributional shifts. By fitting a chain-local baseline from the first
`min_baseline_steps` steps, we detect when the agent's behavioral quality
degrades during the same task.

Algorithm:
  1. Compute step_risk[k] for each step k (lightweight behavioral proxy).
  2. fit_baseline(step_risks[:min_baseline_steps]) — chain-local CUSUM init.
  3. update(step_risk[k]) for k >= min_baseline_steps.
  4. Return DriftResult with first drift point detected.

Parameters: cusum_k=0.1, cusum_h=1.0 (absolute threshold, not std-scaled).

Usage
-----
    from llm_guard.monitoring.intra_chain_drift_detector import IntraChainDriftDetector

    det = IntraChainDriftDetector()
    result = det.detect(chain["steps"])
    if result.drift_detected:
        print(f"Drift at step {result.drift_point}")
Purpose:    Fits a chain-local CUSUM baseline from the first N steps and detects when behavioral quality degrades within the same task. Designed for 15-200 step long-horizon agents.
Exports:    IntraChainDriftDetector
Added:      v0.1.0
Status:     Experimental
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List, Optional

import numpy as np


@dataclass
class DriftResult:
    """Result of drift detection on a single chain."""
    drift_detected: bool
    drift_point: Optional[int]   # step index of first drift alert (None if no drift)
    drift_score: float           # CUSUM value at drift point (or final value if no drift)
    step_risks: List[float]      # per-step risk values
    n_steps: int                 # total number of steps processed


class IntraChainDriftDetector:
    """
    Chain-local drift detector using CUSUM with per-chain baseline.

    Parameters
    ----------
    cusum_k : float
        CUSUM allowance. Should be smaller than the expected risk shift.
        With step_risk in [0, 1], degraded steps score ~0.6 and normal ~0.0,
        so k=0.1 gives a net increment of 0.5 per degraded step. Default: 0.1.
    cusum_h : float
        CUSUM decision threshold (absolute, not std-scaled). Default: 1.0.
    min_baseline_steps : int
        Number of initial steps used to fit the chain-local baseline.
        Default: 3.
    """

    def __init__(
        self,
        cusum_k: float = 0.1,
        cusum_h: float = 1.0,
        min_baseline_steps: int = 3,
    ) -> None:
        self.cusum_k = cusum_k
        self.cusum_h = cusum_h
        self.min_baseline_steps = min_baseline_steps

    @staticmethod
    def _step_risk(step: Dict) -> float:
        """
        Lightweight per-step behavioral risk proxy.

        High risk signals:
          - Empty or very short observation (< 15 chars): retrieval failure
          - Observation contains error keywords: tool failure
          - Short thought (< 10 chars): uncertain reasoning

        Returns float in [0.0, 1.0].
        """
        obs = step.get("observation", "")
        thought = step.get("thought", "")

        risk = 0.0

        # Empty observation = retrieval failure
        if len(obs.strip()) < 15:
            risk += 0.6

        # Error keywords in observation
        obs_lower = obs.lower()
        if any(kw in obs_lower for kw in ("error", "not found", "no result", "failed", "exception")):
            risk += 0.3

        # Very short thought = uncertain reasoning
        if len(thought.strip()) < 10:
            risk += 0.1

        return float(min(1.0, risk))

    def detect(self, steps: List[Dict]) -> DriftResult:
        """
        Detect intra-chain drift over the provided steps.

        Uses CUSUM with per-chain baseline mean. The threshold H is absolute
        (not scaled by baseline std) to remain stable with short baselines.

        Parameters
        ----------
        steps : List[Dict]
            Chain steps with keys: thought, action_type, action_arg, observation.

        Returns
        -------
        DriftResult
        """
        n = len(steps)
        step_risks = [self._step_risk(s) for s in steps]

        if n < self.min_baseline_steps:
            return DriftResult(
                drift_detected=False,
                drift_point=None,
                drift_score=0.0,
                step_risks=step_risks,
                n_steps=n,
            )

        # Baseline statistics from first min_baseline_steps steps
        baseline = step_risks[:self.min_baseline_steps]
        mu = float(np.mean(baseline))

        # CUSUM with absolute threshold
        # H = cusum_h (absolute, not std-relative)
        # k = cusum_k (allowance for expected normal variation)
        cusum_s = 0.0
        drift_point: Optional[int] = None
        drift_score: float = 0.0

        for idx in range(self.min_baseline_steps, n):
            cusum_s = max(0.0, cusum_s + (step_risks[idx] - mu - self.cusum_k))
            drift_score = cusum_s
            if cusum_s > self.cusum_h and drift_point is None:
                drift_point = idx
                break  # stop at first detection

        return DriftResult(
            drift_detected=drift_point is not None,
            drift_point=drift_point,
            drift_score=float(drift_score),
            step_risks=step_risks,
            n_steps=n,
        )
