"""
Token Budget Manager — global resource accounting and cost control.

Design notes:
  - Tracks token usage and USD cost across all agents and tasks
  - Enforces budget limits (hard cap and soft warning)
  - Supports degradation strategies when budget is running low:
    * Switch to a cheaper model
    * Reduce max_iterations
    * Skip optional steps
  - Thread-safe via asyncio locks
"""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Callable
from enum import StrEnum
from typing import Any

from pydantic import BaseModel, Field

from nexus.types import TokenUsage

logger = logging.getLogger("nexus.runtime.budget")


class BudgetStatus(StrEnum):
    """Current budget health."""
    HEALTHY = "healthy"        # < 70% consumed
    WARNING = "warning"        # 70-90% consumed
    CRITICAL = "critical"      # 90-100% consumed
    EXHAUSTED = "exhausted"    # >= 100% consumed


class BudgetConfig(BaseModel):
    """Configuration for the token budget manager."""

    max_total_tokens: int | None = None
    max_total_cost_usd: float | None = None
    warning_threshold: float = 0.7    # 70%
    critical_threshold: float = 0.9   # 90%

    # Degradation strategies
    on_warning: str = "log"           # "log" | "downgrade_model" | "reduce_iterations"
    on_critical: str = "downgrade_model"
    on_exhausted: str = "stop"        # "stop" | "allow_overflow"

    # Model downgrade path
    model_priority: list[str] = Field(
        default_factory=lambda: [
            "gpt-4o",
            "gpt-4o-mini",
            "gpt-3.5-turbo",
        ]
    )


class TokenBudgetManager:
    """
    Global token and cost budget manager.

    Tracks cumulative usage across all agents and enforces limits.

    Usage:
        budget = TokenBudgetManager(
            config=BudgetConfig(max_total_cost_usd=1.0)
        )

        # Check before making an LLM call
        if budget.can_spend(estimated_tokens=1000):
            result = await model.chat(...)
            budget.record(usage)
        else:
            # Handle budget exhaustion
            ...

        # Check budget health
        print(budget.status)  # BudgetStatus.HEALTHY
        print(budget.summary())
    """

    def __init__(self, config: BudgetConfig | None = None) -> None:
        self.config = config or BudgetConfig()
        self._total_usage = TokenUsage()
        self._agent_usage: dict[str, TokenUsage] = {}
        self._task_usage: dict[str, TokenUsage] = {}
        self._call_count = 0
        self._lock = asyncio.Lock()
        self._callbacks: list[Callable[[BudgetStatus], Any]] = []

    @property
    def total_usage(self) -> TokenUsage:
        return self._total_usage

    @property
    def status(self) -> BudgetStatus:
        """Current budget health status."""
        ratio = self._consumption_ratio()
        if ratio >= 1.0:
            return BudgetStatus.EXHAUSTED
        elif ratio >= self.config.critical_threshold:
            return BudgetStatus.CRITICAL
        elif ratio >= self.config.warning_threshold:
            return BudgetStatus.WARNING
        return BudgetStatus.HEALTHY

    def _consumption_ratio(self) -> float:
        """How much of the budget has been consumed (0.0 to 1.0+)."""
        ratios = []
        if self.config.max_total_tokens and self.config.max_total_tokens > 0:
            ratios.append(
                self._total_usage.total_tokens / self.config.max_total_tokens
            )
        if self.config.max_total_cost_usd and self.config.max_total_cost_usd > 0:
            ratios.append(
                self._total_usage.cost_usd / self.config.max_total_cost_usd
            )
        return max(ratios) if ratios else 0.0

    def can_spend(self, estimated_tokens: int = 0, estimated_cost: float = 0.0) -> bool:
        """Check if spending the estimated amount would exceed the budget."""
        if self.status == BudgetStatus.EXHAUSTED:
            return self.config.on_exhausted == "allow_overflow"

        if self.config.max_total_tokens:
            if (self._total_usage.total_tokens + estimated_tokens) > self.config.max_total_tokens:
                return False

        if self.config.max_total_cost_usd:
            if (self._total_usage.cost_usd + estimated_cost) > self.config.max_total_cost_usd:
                return False

        return True

    async def record(
        self,
        usage: TokenUsage,
        agent_id: str | None = None,
        task_id: str | None = None,
    ) -> BudgetStatus:
        """
        Record token usage and return the current budget status.

        Thread-safe via asyncio lock.
        """
        async with self._lock:
            self._total_usage = self._total_usage + usage
            self._call_count += 1

            if agent_id:
                prev = self._agent_usage.get(agent_id, TokenUsage())
                self._agent_usage[agent_id] = prev + usage

            if task_id:
                prev = self._task_usage.get(task_id, TokenUsage())
                self._task_usage[task_id] = prev + usage

        new_status = self.status

        # Fire callbacks on status changes
        for callback in self._callbacks:
            try:
                callback(new_status)
            except Exception as exc:
                logger.error(f"Budget callback error: {exc}")

        if new_status == BudgetStatus.WARNING:
            logger.warning(
                f"Budget warning: {self._consumption_ratio():.0%} consumed "
                f"(tokens={self._total_usage.total_tokens}, "
                f"cost=${self._total_usage.cost_usd:.4f})"
            )
        elif new_status == BudgetStatus.CRITICAL:
            logger.warning(
                f"Budget CRITICAL: {self._consumption_ratio():.0%} consumed!"
            )
        elif new_status == BudgetStatus.EXHAUSTED:
            logger.error(
                f"Budget EXHAUSTED: {self._consumption_ratio():.0%} consumed!"
            )

        return new_status

    def get_recommended_model(self) -> str | None:
        """
        Based on current budget health, recommend which model to use.

        Returns a cheaper model when budget is under pressure.
        """
        models = self.config.model_priority
        if not models:
            return None

        status = self.status
        if status == BudgetStatus.HEALTHY:
            return models[0]  # Use the best model
        elif status == BudgetStatus.WARNING:
            return models[min(1, len(models) - 1)]  # Downgrade one level
        elif status == BudgetStatus.CRITICAL:
            return models[-1]  # Use the cheapest model
        else:
            return None  # Budget exhausted

    def on_status_change(self, callback: Callable[[BudgetStatus], Any]) -> None:
        """Register a callback for budget status changes."""
        self._callbacks.append(callback)

    def get_agent_usage(self, agent_id: str) -> TokenUsage:
        """Get cumulative usage for a specific agent."""
        return self._agent_usage.get(agent_id, TokenUsage())

    def get_task_usage(self, task_id: str) -> TokenUsage:
        """Get cumulative usage for a specific task."""
        return self._task_usage.get(task_id, TokenUsage())

    def summary(self) -> dict[str, Any]:
        """Generate a human-readable summary of budget status."""
        return {
            "status": self.status.value,
            "consumption": f"{self._consumption_ratio():.1%}",
            "total_tokens": self._total_usage.total_tokens,
            "total_cost_usd": round(self._total_usage.cost_usd, 4),
            "call_count": self._call_count,
            "limits": {
                "max_tokens": self.config.max_total_tokens,
                "max_cost_usd": self.config.max_total_cost_usd,
            },
            "agents_tracked": len(self._agent_usage),
            "tasks_tracked": len(self._task_usage),
        }

    def reset(self) -> None:
        """Reset all usage counters."""
        self._total_usage = TokenUsage()
        self._agent_usage.clear()
        self._task_usage.clear()
        self._call_count = 0
