"""VERONICA Budget Enforcement - Chain-level cost ceiling for LLM calls."""
# nogil-audited: 2026-03-08
# Findings:
#   - All shared mutable fields (_spent_usd, _call_count, _exceeded) are
#     accessed exclusively inside ``with self._lock:``. No changes needed.

from __future__ import annotations

__all__ = ["BudgetEnforcer"]
from dataclasses import dataclass, field
from typing import Dict
import logging
import math
import threading

from veronica_core.kernel.decision import ReasonCode, make_envelope
from veronica_core.runtime_policy import PolicyContext, PolicyDecision

logger = logging.getLogger(__name__)


@dataclass
class BudgetEnforcer:
    """Chain-level budget enforcement for LLM API calls.

    Tracks cumulative cost across a chain of calls and stops
    when the budget ceiling is reached. Thread-safe.

    Example:
        budget = BudgetEnforcer(limit_usd=100.0)
        for call in llm_calls:
            if not budget.spend(call.cost):
                break  # Budget exceeded
    """

    limit_usd: float = 100.0
    _spent_usd: float = field(default=0.0, init=False)

    def __post_init__(self) -> None:
        if math.isnan(self.limit_usd) or math.isinf(self.limit_usd):
            raise ValueError(f"limit_usd must be finite, got {self.limit_usd!r}")
        if self.limit_usd < 0:
            raise ValueError(f"limit_usd must be non-negative, got {self.limit_usd!r}")
        # Zero budget means "no spending allowed" -- mark as exceeded immediately
        # so is_exceeded is consistent with check()/spend() from construction.
        if self.limit_usd == 0.0:
            self._exceeded = True

    _call_count: int = field(default=0, init=False)
    _exceeded: bool = field(default=False, init=False)
    _lock: threading.Lock = field(
        default_factory=threading.Lock, init=False, repr=False
    )

    def spend(self, amount_usd: float) -> bool:
        """Record spending. Returns True if within budget, False if exceeded.

        Args:
            amount_usd: Cost of this call in USD

        Returns:
            True if still within budget after this spend

        Raises:
            ValueError: If amount_usd is negative.
        """
        with self._lock:
            if math.isnan(amount_usd) or math.isinf(amount_usd):
                raise ValueError(f"amount must be a finite number, got {amount_usd}")
            if amount_usd < 0:
                raise ValueError(f"amount must be non-negative, got {amount_usd}")
            # Zero-budget: deny all spending (consistent with check()).
            if self.limit_usd == 0.0:
                if not self._exceeded:
                    logger.warning(
                        "[VERONICA_BUDGET] Budget is zero: no spending permitted"
                    )
                    self._exceeded = True
                return False
            projected = self._spent_usd + amount_usd
            if projected > self.limit_usd:
                logger.warning(
                    f"[VERONICA_BUDGET] Budget exceeded: "
                    f"${projected:.2f} / ${self.limit_usd:.2f} "
                    f"({self._call_count} calls)"
                )
                self._exceeded = True
                return False
            self._spent_usd = projected
            self._call_count += 1
            return True

    @property
    def spent_usd(self) -> float:
        """Total amount spent so far."""
        with self._lock:
            return self._spent_usd

    @property
    def remaining_usd(self) -> float:
        """Remaining budget in USD.

        Returns 0.0 when the budget has been exceeded (is_exceeded is True),
        even if limit_usd - spent_usd would be negative. This means the return
        value is always in the range [0.0, limit_usd] and can be used safely
        for progress displays without sign checks.

        Note: a return value of 0.0 does NOT distinguish between "exactly at
        the limit" and "exceeded". Use is_exceeded for that distinction.
        """
        with self._lock:
            if self._exceeded:
                return 0.0
            return max(0.0, self.limit_usd - self._spent_usd)

    @property
    def is_exceeded(self) -> bool:
        """True if a spend attempt has been rejected due to the budget limit."""
        with self._lock:
            return self._exceeded

    @property
    def call_count(self) -> int:
        """Total number of calls tracked."""
        with self._lock:
            return self._call_count

    @property
    def utilization(self) -> float:
        """Budget utilization as a fraction (0.0 to 1.0+).

        Returns 1.0 (100% utilized) when ``limit_usd == 0`` -- a zero budget
        is immediately exhausted by definition, not infinitely exceeded.
        """
        with self._lock:
            if self.limit_usd == 0.0:
                return 1.0  # 100% utilized: zero-budget is always exhausted
            return self._spent_usd / self.limit_usd

    def reset(self) -> None:
        """Reset budget tracking."""
        with self._lock:
            self._spent_usd = 0.0
            self._call_count = 0
            # Re-apply zero-budget invariant: zero budget is always exceeded.
            self._exceeded = self.limit_usd == 0.0
            logger.info("[VERONICA_BUDGET] Budget reset")

    @property
    def policy_type(self) -> str:
        """RuntimePolicy protocol: policy type identifier."""
        return "budget"

    def check(self, context: PolicyContext) -> PolicyDecision:
        """RuntimePolicy protocol: check if operation is within budget.

        Evaluates whether the projected cost (spent + context.cost_usd)
        would exceed the budget limit. Does NOT record spending --
        use spend() to record actual costs after the operation.

        Args:
            context: PolicyContext with cost_usd set to projected cost

        Returns:
            PolicyDecision allowing or denying the operation
        """
        with self._lock:
            # Zero-budget: deny all calls regardless of cost_estimate.
            # A limit of 0.0 means no spending is permitted at all.
            if self.limit_usd == 0.0:
                reason = "Budget is zero: no spending permitted"
                return PolicyDecision(
                    allowed=False,
                    policy_type=self.policy_type,
                    reason=reason,
                    envelope=make_envelope(
                        decision="DENY",
                        reason_code=ReasonCode.BUDGET_EXCEEDED,
                        reason=reason,
                        issuer="BudgetEnforcer",
                    ),
                )
            if self._exceeded:
                reason = (
                    f"Budget exceeded: ${self._spent_usd:.2f} / ${self.limit_usd:.2f}"
                )
                return PolicyDecision(
                    allowed=False,
                    policy_type=self.policy_type,
                    reason=reason,
                    envelope=make_envelope(
                        decision="DENY",
                        reason_code=ReasonCode.BUDGET_EXCEEDED,
                        reason=reason,
                        issuer="BudgetEnforcer",
                    ),
                )
            # Validate cost_usd: NaN compares as False in > checks, letting
            # invalid costs through silently; negative values reduce projected
            # spend and bypass limits near the ceiling.
            cost = context.cost_usd
            if math.isnan(cost) or math.isinf(cost) or cost < 0:
                reason = f"Invalid cost_usd in context: {cost!r}"
                return PolicyDecision(
                    allowed=False,
                    policy_type=self.policy_type,
                    reason=reason,
                    envelope=make_envelope(
                        decision="DENY",
                        reason_code=ReasonCode.UNKNOWN,
                        reason=reason,
                        issuer="BudgetEnforcer",
                    ),
                )
            projected = self._spent_usd + cost
            if projected > self.limit_usd:
                reason = (
                    f"Budget would exceed: ${projected:.2f} > ${self.limit_usd:.2f}"
                )
                return PolicyDecision(
                    allowed=False,
                    policy_type=self.policy_type,
                    reason=reason,
                    envelope=make_envelope(
                        decision="DENY",
                        reason_code=ReasonCode.BUDGET_EXCEEDED,
                        reason=reason,
                        issuer="BudgetEnforcer",
                    ),
                )
            return PolicyDecision(allowed=True, policy_type=self.policy_type)

    def to_dict(self) -> Dict:
        """Serialize budget state."""
        with self._lock:
            return {
                "limit_usd": self.limit_usd,
                "spent_usd": self._spent_usd,
                "call_count": self._call_count,
            }
