"""Budget controller — enforces hard execution limits before every model call."""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from datetime import UTC, datetime

from dpe.models.config import BudgetLimits, ThrottleBehavior
from dpe.models.manifest import ThrottleEvent

logger = logging.getLogger(__name__)


class BudgetExceededError(Exception):
    """Raised when a hard budget limit is breached."""

    def __init__(self, reason: str) -> None:
        self.reason = reason
        super().__init__(reason)


@dataclass
class CallDecision:
    """Result of a pre-call budget check."""

    allowed: bool
    degraded: bool = False
    reason: str = ""


@dataclass
class BudgetController:
    """Central enforcement point for all execution limits.

    Before every model call, the orchestrator must call pre_call_check().
    If any limit is breached, the call is denied.
    """

    limits: BudgetLimits
    total_calls_used: int = 0
    premium_calls_used: int = 0
    critique_rounds_used: int = 0
    start_time: float = field(default_factory=time.monotonic)
    stage_attempts: dict[str, int] = field(default_factory=dict)
    throttle_events: list[ThrottleEvent] = field(default_factory=list)
    _degraded: bool = False

    @property
    def wall_time_elapsed(self) -> float:
        return time.monotonic() - self.start_time

    @property
    def is_degraded(self) -> bool:
        return self._degraded

    def pre_call_check(self, *, is_premium: bool, stage: str, adapter: str) -> CallDecision:
        """Gate check before every model call. Returns a CallDecision."""
        self.stage_attempts[stage] = self.stage_attempts.get(stage, 0) + 1

        # Wall-time check
        if self.wall_time_elapsed >= self.limits.max_wall_time_seconds:
            return self._handle_breach(
                f"Wall time exceeded: {self.wall_time_elapsed:.0f}s"
                f" >= {self.limits.max_wall_time_seconds}s",
                stage=stage,
                adapter=adapter,
            )

        # Total calls check
        if self.total_calls_used >= self.limits.max_total_calls:
            return self._handle_breach(
                f"Total calls exceeded: {self.total_calls_used} >= {self.limits.max_total_calls}",
                stage=stage,
                adapter=adapter,
            )

        # Premium calls check
        if is_premium and self.premium_calls_used >= self.limits.max_premium_calls:
            return self._handle_breach(
                f"Premium calls exceeded: {self.premium_calls_used}"
                f" >= {self.limits.max_premium_calls}",
                stage=stage,
                adapter=adapter,
            )

        return CallDecision(allowed=True, degraded=self._degraded)

    def record_call(self, *, is_premium: bool) -> None:
        """Record that a model call was made. Counters are monotonically increasing."""
        self.total_calls_used += 1
        if is_premium:
            self.premium_calls_used += 1
        logger.info(
            "Budget: total=%d/%d premium=%d/%d wall=%.0fs/%ds",
            self.total_calls_used,
            self.limits.max_total_calls,
            self.premium_calls_used,
            self.limits.max_premium_calls,
            self.wall_time_elapsed,
            self.limits.max_wall_time_seconds,
        )

    def record_critique_round(self) -> None:
        """Record a completed critique round."""
        self.critique_rounds_used += 1

    def can_critique(self) -> bool:
        """Check if another critique round is allowed."""
        return self.critique_rounds_used < self.limits.max_critique_rounds

    def _handle_breach(self, reason: str, *, stage: str, adapter: str) -> CallDecision:
        """Handle a budget limit breach based on throttle_behavior."""
        logger.warning("Budget breach: %s", reason)

        event = ThrottleEvent(
            timestamp=datetime.now(UTC),
            adapter=adapter,
            stage=stage,
            action_taken=self.limits.throttle_behavior.value,
        )
        self.throttle_events.append(event)

        if self.limits.throttle_behavior == ThrottleBehavior.ABORT:
            return CallDecision(allowed=False, reason=reason)

        if self.limits.throttle_behavior == ThrottleBehavior.DEGRADE_TO_LOCAL:
            self._degraded = True
            return CallDecision(allowed=True, degraded=True, reason=reason)

        if self.limits.throttle_behavior == ThrottleBehavior.RETRY_THEN_DEGRADE:
            if not self._degraded:
                self._degraded = True
                return CallDecision(allowed=True, degraded=False, reason=reason)
            return CallDecision(allowed=True, degraded=True, reason=reason)

        return CallDecision(allowed=False, reason=reason)

    def snapshot(self) -> dict[str, object]:
        """Return current budget state for manifest."""
        return {
            "total_calls_used": self.total_calls_used,
            "premium_calls_used": self.premium_calls_used,
            "critique_rounds_used": self.critique_rounds_used,
            "wall_time_seconds": round(self.wall_time_elapsed, 2),
            "degraded": self._degraded,
            "throttle_events": len(self.throttle_events),
        }
