"""
ConformalRiskBudget — per-step risk allocation for multi-step agent chains.

Gives downstream systems a principled way to bound total chain failure
probability.  Three allocation methods are supported:

  bonferroni : per_step_threshold = total_budget / n_steps  (union-bound)
  product    : per_step_threshold = 1-(1-total_budget)^(1/n_steps)
               (exact under step-independence assumption)
  adaptive   : redistribute the *remaining* budget equally across remaining
               steps after each observation — tightest early-abort signal.

Usage
-----
    budget = ConformalRiskBudget(total_budget=0.20)
    thresholds = budget.allocate(n_steps=5)

    for i, step_risk in enumerate(chain_step_risks):
        abort = budget.observe(i, step_risk)
        if abort:
            handle_abort()

    print(budget.summary())
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Literal


@dataclass
class BudgetState:
    step_idx: int
    step_risk: float
    allocated_threshold: float
    cumulative_used: float
    remaining: float
    over_budget: bool


class ConformalRiskBudget:
    def __init__(
        self,
        total_budget: float = 0.20,
        method: Literal["bonferroni", "product", "adaptive"] = "product",
    ) -> None:
        if not 0 < total_budget < 1:
            raise ValueError("total_budget must be in (0, 1)")
        self._budget = total_budget
        self._method = method
        self._n_steps: int = 0
        self._thresholds: List[float] = []
        self._history: List[BudgetState] = []
        self._cumulative: float = 0.0  # union-bound accumulator

    # ------------------------------------------------------------------
    def allocate(self, n_steps: int) -> List[float]:
        """Return per-step abort thresholds for a chain of `n_steps` steps."""
        if n_steps <= 0:
            raise ValueError("n_steps must be >= 1")
        self._n_steps = n_steps
        self._thresholds = self._compute(n_steps, self._budget)
        return list(self._thresholds)

    def _compute(self, n: int, budget: float) -> List[float]:
        if self._method == "bonferroni":
            t = budget / n
            return [t] * n
        if self._method == "product":
            t = 1.0 - (1.0 - budget) ** (1.0 / n)
            return [t] * n
        # adaptive: equal share of remaining budget, recomputed at observe()
        t = budget / n
        return [t] * n

    # ------------------------------------------------------------------
    def observe(self, step_idx: int, risk: float) -> bool:
        """Record a step's risk score; returns True if chain should abort."""
        if not self._thresholds:
            raise RuntimeError("Call allocate() before observe()")
        if self._method == "adaptive":
            steps_left = self._n_steps - step_idx
            remaining_budget = max(0.0, self._budget - self._cumulative)
            threshold = remaining_budget / max(steps_left, 1)
        else:
            threshold = self._thresholds[min(step_idx, len(self._thresholds) - 1)]

        self._cumulative = min(1.0, self._cumulative + risk)
        over = self._cumulative > self._budget

        self._history.append(BudgetState(
            step_idx=step_idx,
            step_risk=risk,
            allocated_threshold=threshold,
            cumulative_used=self._cumulative,
            remaining=max(0.0, self._budget - self._cumulative),
            over_budget=over,
        ))
        return over

    # ------------------------------------------------------------------
    @property
    def remaining_budget(self) -> float:
        return max(0.0, self._budget - self._cumulative)

    def reset(self) -> None:
        self._history.clear()
        self._cumulative = 0.0

    def summary(self) -> dict:
        return {
            "total_budget": self._budget,
            "method": self._method,
            "n_steps": self._n_steps,
            "cumulative_used": round(self._cumulative, 4),
            "remaining": round(self.remaining_budget, 4),
            "over_budget": self._cumulative > self._budget,
            "n_observed": len(self._history),
            "aborted_at": next(
                (s.step_idx for s in self._history if s.over_budget), None
            ),
        }
