"""
FailureRateAdaptiveSampler — rolling failure-rate adaptive guard sampling.

Reduces guard overhead on stable pipelines by sampling a fraction of calls
when the pipeline is healthy, and bursting to 100% when failure rate spikes.

Standalone: no dependencies on AgentGuard or any API client. Works with any
object that has ``needs_alert: bool`` and ``confidence_tier: str`` attributes.
Purpose:    Samples a reduced fraction of calls when the pipeline is healthy and bursts to 100% sampling when failure rate spikes. Reduces guard overhead without missing failure clusters.
Exports:    FailureRateAdaptiveSampler
Added:      v0.1.0
Status:     Stable
"""
from __future__ import annotations

import random
from collections import deque
from typing import Dict


class FailureRateAdaptiveSampler:
    """
    Controls guard call rate based on a rolling failure-rate estimate.

    State machine per pipeline_id:

        cold_start (n_observations < window)
            → should_score() always True
            → transitions to healthy or burst once n_observations == window

        healthy (failure_rate < anomaly_threshold)
            → should_score() returns True with probability healthy_sample_rate
            → transitions to burst if failure_rate >= anomaly_threshold

        burst (failure_rate >= anomaly_threshold)
            → should_score() always True
            → burst_counter decrements on each should_score() call
            → when burst_counter reaches 0: if rate dropped → healthy, else reset burst

    Parameters
    ----------
    window : int, default 50
        Rolling window size for failure rate estimation.
    healthy_sample_rate : float, default 0.15
        Fraction of calls to score when healthy (0.15 = 15%).
    anomaly_threshold : float, default 0.25
        Failure rate above which pipeline enters burst mode.
    burst_window : int, default 20
        Number of should_score() calls to run at 100% after anomaly detected.
    seed : int, default 42
        RNG seed for reproducible sampling decisions.
    """

    def __init__(
        self,
        window: int = 50,
        healthy_sample_rate: float = 0.15,
        anomaly_threshold: float = 0.25,
        burst_window: int = 20,
        seed: int = 42,
    ) -> None:
        self._window = window
        self._healthy_rate = healthy_sample_rate
        self._anomaly_threshold = anomaly_threshold
        self._burst_window = burst_window
        self._rng = random.Random(seed)

        # Per-pipeline state
        self._buffers: Dict[str, deque] = {}       # rolling 0/1 outcomes
        self._modes: Dict[str, str] = {}           # "cold_start"|"healthy"|"burst"
        self._burst_counters: Dict[str, int] = {}  # calls remaining in burst
        self._n_obs: Dict[str, int] = {}           # total observe() calls

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def should_score(self, pipeline_id: str = "default") -> bool:
        """
        Returns True if this call should be sent through the guard.

        In cold_start mode: always True.
        In burst mode: always True; decrements the burst counter by 1.
        In healthy mode: True with probability healthy_sample_rate.

        The burst counter decrements on EVERY call to should_score() during burst,
        regardless of whether observe() is subsequently called. This keeps the
        burst window measured in wall-clock calls to the sampler.
        """
        self._ensure_pipeline(pipeline_id)
        mode = self._modes[pipeline_id]

        if mode == "cold_start":
            return True

        if mode == "burst":
            # Decrement burst counter; this call is always scored
            self._burst_counters[pipeline_id] -= 1
            return True

        # healthy mode
        return self._rng.random() < self._healthy_rate

    def observe(self, result: object, pipeline_id: str = "default") -> None:
        """
        Record a guard result. A chain is a failure if:
            result.needs_alert is True  OR  result.confidence_tier == "LOW"

        Safe to call even if should_score() was not called first — the observation
        is recorded in the rolling window without affecting the burst counter.
        If the guard call fails (network error), do NOT call observe().
        """
        self._ensure_pipeline(pipeline_id)
        is_fail = getattr(result, "needs_alert", False) or \
                  getattr(result, "confidence_tier", "") == "LOW"
        self._buffers[pipeline_id].append(1 if is_fail else 0)
        self._n_obs[pipeline_id] += 1

        # Transitions driven by observations
        mode = self._modes[pipeline_id]
        fr = self._failure_rate(pipeline_id)

        if mode == "cold_start":
            if self._n_obs[pipeline_id] >= self._window:
                if fr >= self._anomaly_threshold:
                    self._modes[pipeline_id] = "burst"
                    self._burst_counters[pipeline_id] = self._burst_window
                else:
                    self._modes[pipeline_id] = "healthy"

        elif mode == "healthy":
            if fr >= self._anomaly_threshold:
                self._modes[pipeline_id] = "burst"
                self._burst_counters[pipeline_id] = self._burst_window

        elif mode == "burst":
            # Re-evaluate once the burst window's should_score() calls are exhausted
            if self._burst_counters[pipeline_id] <= 0:
                if fr < self._anomaly_threshold:
                    self._modes[pipeline_id] = "healthy"
                else:
                    # Still unhealthy: reset burst counter for another window
                    self._burst_counters[pipeline_id] = self._burst_window

    def status(self, pipeline_id: str = "default") -> dict:
        """
        Returns a diagnostic snapshot for the given pipeline:
        {
            "pipeline_id":      str,
            "failure_rate":     float,
            "mode":             "cold_start" | "healthy" | "burst",
            "calls_in_burst":   int,   # should_score() calls remaining (0 if not burst)
            "n_observations":   int,
            "sample_rate":      float, # current effective rate
        }
        """
        self._ensure_pipeline(pipeline_id)
        mode = self._modes[pipeline_id]
        rate = 1.0 if mode in ("cold_start", "burst") else self._healthy_rate
        return {
            "pipeline_id": pipeline_id,
            "failure_rate": self._failure_rate(pipeline_id),
            "mode": mode,
            "calls_in_burst": max(0, self._burst_counters[pipeline_id]),
            "n_observations": self._n_obs[pipeline_id],
            "sample_rate": rate,
        }

    def reset(self, pipeline_id: str = "default") -> None:
        """Clear all state for a pipeline (e.g., after a deployment)."""
        for d in (self._buffers, self._modes, self._burst_counters, self._n_obs):
            d.pop(pipeline_id, None)

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _ensure_pipeline(self, pipeline_id: str) -> None:
        if pipeline_id not in self._buffers:
            buf = deque(maxlen=self._window)
            # Pre-fill with zeros so failure_rate is well-defined from first observe()
            buf.extend([0] * self._window)
            self._buffers[pipeline_id] = buf
            self._modes[pipeline_id] = "cold_start"
            self._burst_counters[pipeline_id] = 0
            self._n_obs[pipeline_id] = 0

    def _failure_rate(self, pipeline_id: str) -> float:
        buf = self._buffers.get(pipeline_id)
        if not buf:
            return 0.0
        return sum(buf) / len(buf)
