"""Budget engine — enforce spending limits and track LLM costs."""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from typing import Any

from agent_scout.events import Event, EventBus, EventType

logger = logging.getLogger(__name__)


@dataclass
class BudgetConfig:
    """Budget limits for LLM spending."""

    # Per-task hard limit (USD)
    task_limit: float = 1.0
    # Daily hard limit (USD)
    daily_limit: float = 10.0
    # Monthly hard limit (USD)
    monthly_limit: float = 50.0
    # Whether to enforce limits (False = track only)
    enforce: bool = True


@dataclass
class CostRecord:
    """A single LLM call cost record."""

    model: str
    input_tokens: int
    output_tokens: int
    cost_usd: float
    timestamp: float = field(default_factory=time.time)
    task_id: str | None = None


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

    def __init__(self, limit_type: str, limit: float, current: float) -> None:
        self.limit_type = limit_type
        self.limit = limit
        self.current = current
        super().__init__(f"{limit_type} budget exceeded: ${current:.4f} >= ${limit:.4f}")


class BudgetEngine:
    """Track LLM costs and enforce spending limits.

    Listens to COST_UPDATE events from LLMClient and maintains running
    totals per task, per day, and per month.
    """

    def __init__(
        self,
        config: BudgetConfig | None = None,
        event_bus: EventBus | None = None,
    ) -> None:
        self._config = config or BudgetConfig()
        self._bus = event_bus
        self._records: list[CostRecord] = []

        # Subscribe to cost events if bus is provided
        if self._bus:
            self._bus.subscribe(self._on_cost_event, EventType.COST_UPDATE)

    @property
    def config(self) -> BudgetConfig:
        return self._config

    def record_cost(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        cost_usd: float,
        task_id: str | None = None,
    ) -> None:
        """Record a cost and check limits."""
        record = CostRecord(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost_usd=cost_usd,
            task_id=task_id,
        )
        self._records.append(record)

    async def check_budget(self, task_id: str | None = None) -> None:
        """Check if any budget limit has been exceeded.

        Raises BudgetExceededError if enforce is True and a limit is hit.
        Emits BUDGET_WARNING at 80% and BUDGET_EXCEEDED at 100%.
        """
        now = time.time()

        # Task budget
        if task_id:
            task_total = self.get_task_cost(task_id)
            await self._check_limit("task", task_total, self._config.task_limit, task_id)

        # Daily budget
        daily_total = self.get_daily_cost(now)
        await self._check_limit("daily", daily_total, self._config.daily_limit)

        # Monthly budget
        monthly_total = self.get_monthly_cost(now)
        await self._check_limit("monthly", monthly_total, self._config.monthly_limit)

    def get_task_cost(self, task_id: str) -> float:
        """Total cost for a specific task."""
        return sum(r.cost_usd for r in self._records if r.task_id == task_id)

    def get_daily_cost(self, now: float | None = None) -> float:
        """Total cost for the current day (last 24 hours)."""
        now = now or time.time()
        day_start = now - 86400
        return sum(r.cost_usd for r in self._records if r.timestamp >= day_start)

    def get_monthly_cost(self, now: float | None = None) -> float:
        """Total cost for the current month (last 30 days)."""
        now = now or time.time()
        month_start = now - 86400 * 30
        return sum(r.cost_usd for r in self._records if r.timestamp >= month_start)

    def get_total_tokens(self, task_id: str | None = None) -> dict[str, int]:
        """Get total token counts, optionally filtered by task."""
        records = self._records
        if task_id:
            records = [r for r in records if r.task_id == task_id]
        return {
            "input_tokens": sum(r.input_tokens for r in records),
            "output_tokens": sum(r.output_tokens for r in records),
        }

    def get_summary(self, task_id: str | None = None) -> dict[str, Any]:
        """Get a cost summary."""
        now = time.time()
        tokens = self.get_total_tokens(task_id)
        return {
            "task_cost": self.get_task_cost(task_id) if task_id else 0,
            "daily_cost": self.get_daily_cost(now),
            "monthly_cost": self.get_monthly_cost(now),
            "total_calls": len([r for r in self._records if not task_id or r.task_id == task_id]),
            **tokens,
            "limits": {
                "task": self._config.task_limit,
                "daily": self._config.daily_limit,
                "monthly": self._config.monthly_limit,
            },
        }

    def reset(self) -> None:
        """Clear all cost records."""
        self._records.clear()

    async def _check_limit(
        self,
        limit_type: str,
        current: float,
        limit: float,
        task_id: str | None = None,
    ) -> None:
        """Check a single limit, emit warnings, and optionally enforce."""
        ratio = current / limit if limit > 0 else 0

        if ratio >= 1.0:
            if self._bus:
                await self._bus.emit(
                    Event(
                        event_type=EventType.BUDGET_EXCEEDED,
                        data={
                            "limit_type": limit_type,
                            "limit": limit,
                            "current": current,
                        },
                        task_id=task_id,
                    )
                )
            if self._config.enforce:
                raise BudgetExceededError(limit_type, limit, current)

        elif ratio >= 0.8:
            if self._bus:
                await self._bus.emit(
                    Event(
                        event_type=EventType.BUDGET_WARNING,
                        data={
                            "limit_type": limit_type,
                            "limit": limit,
                            "current": current,
                            "percentage": round(ratio * 100, 1),
                        },
                        task_id=task_id,
                    )
                )

    async def _on_cost_event(self, event: Event) -> None:
        """Handle COST_UPDATE events from the LLM client."""
        data = event.data
        model = data.get("model", "unknown")
        input_tokens = data.get("input_tokens", 0)
        output_tokens = data.get("output_tokens", 0)
        cost_usd = data.get("cost_usd", 0.0)
        task_id = event.task_id

        self.record_cost(model, input_tokens, output_tokens, cost_usd, task_id)
