"""Cost tracking and budget enforcement for provider usage.

Tracks token consumption and monetary costs per plan execution,
enforces budget limits, and emits warnings/blocks when thresholds
are approached or exceeded.
"""

from __future__ import annotations

import logging
import threading
from dataclasses import dataclass
from datetime import date, datetime
from enum import StrEnum
from typing import ClassVar

from cleveragents.domain.models.core.cost_metadata import (
    BudgetExhaustionEvent,
    CostMetadata,
)
from cleveragents.providers.cost_table import CostEntry, ProviderCostTable

logger = logging.getLogger(__name__)


class BudgetStatus(StrEnum):
    """Budget enforcement status levels."""

    UNDER_BUDGET = "under_budget"
    WARNING = "warning"
    EXCEEDED = "exceeded"


# Severity ordering for BudgetStatus (higher = worse)
_BUDGET_SEVERITY: dict[BudgetStatus, int] = {
    BudgetStatus.UNDER_BUDGET: 0,
    BudgetStatus.WARNING: 1,
    BudgetStatus.EXCEEDED: 2,
}


@dataclass(frozen=True, slots=True)
class BudgetCheckResult:
    """Result of a budget check operation.

    Attributes:
        status: Current budget status.
        used: Amount used so far (USD).
        limit: Budget limit (USD). None if unlimited.
        remaining: Amount remaining (USD). None if unlimited.
        budget_type: Which budget was checked ('plan' or 'daily').
    """

    status: BudgetStatus
    used: float
    limit: float | None
    remaining: float | None
    budget_type: str


class CostTracker:
    """Track costs and enforce budget limits for provider usage.

    Manages per-plan and per-day budget tracking with configurable
    warning and blocking thresholds. Uses dependency injection for
    the cost table.

    Attributes:
        WARN_THRESHOLD: Fraction of budget at which warnings are emitted.
        BLOCK_THRESHOLD: Fraction of budget at which requests are blocked.
    """

    WARN_THRESHOLD: ClassVar[float] = 0.9
    BLOCK_THRESHOLD: ClassVar[float] = 1.0

    def __init__(
        self,
        *,
        budget_per_plan: float | None = None,
        budget_per_day: float | None = None,
        cost_table: ProviderCostTable | None = None,
    ) -> None:
        """Initialize the cost tracker.

        Args:
            budget_per_plan: Maximum USD spend per plan (None = unlimited).
            budget_per_day: Maximum USD spend per day (None = unlimited).
            cost_table: Provider cost lookup table (uses defaults if None).

        Raises:
            ValueError: If budget values are negative.
        """
        if budget_per_plan is not None and budget_per_plan < 0.0:
            raise ValueError(f"budget_per_plan must be >= 0, got {budget_per_plan}")
        if budget_per_day is not None and budget_per_day < 0.0:
            raise ValueError(f"budget_per_day must be >= 0, got {budget_per_day}")
        self._budget_per_plan = budget_per_plan
        self._budget_per_day = budget_per_day
        self._cost_table = cost_table or ProviderCostTable()
        self._daily_costs: dict[str, float] = {}
        self._daily_costs_lock: threading.Lock = threading.Lock()

    @property
    def budget_per_plan(self) -> float | None:
        """Return the per-plan budget limit."""
        return self._budget_per_plan

    @property
    def budget_per_day(self) -> float | None:
        """Return the per-day budget limit."""
        return self._budget_per_day

    def get_cost_entry(self, provider: str, model: str) -> CostEntry:
        """Look up the cost entry for a provider/model.

        Args:
            provider: Provider name.
            model: Model identifier.

        Returns:
            The cost entry for the provider/model pair.

        Raises:
            ValueError: If provider or model is empty.
        """
        if not provider:
            raise ValueError("provider must be a non-empty string")
        if not model:
            raise ValueError("model must be a non-empty string")
        return self._cost_table.get_cost_entry(provider, model)

    def estimate_cost(
        self,
        provider: str,
        model: str,
        input_tokens: int,
        output_tokens: int,
    ) -> float:
        """Estimate cost for a token usage.

        Args:
            provider: Provider name.
            model: Model identifier.
            input_tokens: Number of input tokens.
            output_tokens: Number of output tokens.

        Returns:
            Estimated cost in USD.

        Raises:
            ValueError: If arguments are invalid.
        """
        if not provider:
            raise ValueError("provider must be a non-empty string")
        if not model:
            raise ValueError("model must be a non-empty string")
        if input_tokens < 0:
            raise ValueError(f"input_tokens must be >= 0, got {input_tokens}")
        if output_tokens < 0:
            raise ValueError(f"output_tokens must be >= 0, got {output_tokens}")
        entry = self._cost_table.get_cost_entry(provider, model)
        return entry.estimate_cost(input_tokens, output_tokens)

    def record_usage(
        self,
        cost_metadata: CostMetadata,
        *,
        provider: str,
        model: str,
        input_tokens: int,
        output_tokens: int,
    ) -> BudgetCheckResult:
        """Record usage and check budget status.

        Updates cost_metadata in-place with token/cost data.
        Also updates internal daily cost tracking.
        Records budget exhaustion events when thresholds are crossed.

        Args:
            cost_metadata: The plan's cost metadata to update.
            provider: Provider name.
            model: Model identifier.
            input_tokens: Input tokens used.
            output_tokens: Output tokens used.

        Returns:
            Budget check result after recording usage.

        Raises:
            ValueError: If arguments are invalid.
        """
        if not provider:
            raise ValueError("provider must be a non-empty string")
        if not model:
            raise ValueError("model must be a non-empty string")
        if input_tokens < 0:
            raise ValueError(f"input_tokens must be >= 0, got {input_tokens}")
        if output_tokens < 0:
            raise ValueError(f"output_tokens must be >= 0, got {output_tokens}")

        cost = self.estimate_cost(provider, model, input_tokens, output_tokens)
        cost_metadata.record_usage(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost=cost,
            provider=provider,
        )

        with self._daily_costs_lock:
            today_key = date.today().isoformat()
            self._daily_costs[today_key] = self._daily_costs.get(today_key, 0.0) + cost

        if self._budget_per_plan is not None:
            cost_metadata.budget_remaining = max(
                0.0, self._budget_per_plan - cost_metadata.total_cost
            )

        result = self._check_budgets(cost_metadata, provider, model)
        return result

    def check_plan_budget(self, cost_metadata: CostMetadata) -> BudgetCheckResult:
        """Check the per-plan budget without recording usage.

        Args:
            cost_metadata: The plan's cost metadata.

        Returns:
            Budget check result for the plan budget.
        """
        return self._evaluate_budget(
            used=cost_metadata.total_cost,
            limit=self._budget_per_plan,
            budget_type="plan",
        )

    def check_daily_budget(self) -> BudgetCheckResult:
        """Check the per-day budget without recording usage.

        Returns:
            Budget check result for the daily budget.
        """
        with self._daily_costs_lock:
            today_key = date.today().isoformat()
            used = self._daily_costs.get(today_key, 0.0)
        return self._evaluate_budget(
            used=used,
            limit=self._budget_per_day,
            budget_type="daily",
        )

    def get_daily_spend(self) -> float:
        """Return total spend for the current day."""
        with self._daily_costs_lock:
            today_key = date.today().isoformat()
            return self._daily_costs.get(today_key, 0.0)

    def _check_budgets(
        self,
        cost_metadata: CostMetadata,
        provider: str,
        model: str,
    ) -> BudgetCheckResult:
        """Check both plan and daily budgets, return worst status."""
        plan_result = self.check_plan_budget(cost_metadata)
        daily_result = self.check_daily_budget()

        worst = plan_result
        if _BUDGET_SEVERITY.get(daily_result.status, 0) > _BUDGET_SEVERITY.get(
            worst.status, 0
        ):
            worst = daily_result

        if plan_result.status == BudgetStatus.WARNING:
            logger.warning(
                "Plan budget warning: %.2f%% used (%.4f / %.4f USD)",
                (plan_result.used / plan_result.limit * 100)
                if plan_result.limit
                else 0,
                plan_result.used,
                plan_result.limit or 0,
            )
        if daily_result.status == BudgetStatus.WARNING:
            logger.warning(
                "Daily budget warning: %.2f%% used (%.4f / %.4f USD)",
                (daily_result.used / daily_result.limit * 100)
                if daily_result.limit
                else 0,
                daily_result.used,
                daily_result.limit or 0,
            )

        if plan_result.status == BudgetStatus.EXCEEDED:
            self._record_exhaustion_event(cost_metadata, "plan", provider, model)
        if daily_result.status == BudgetStatus.EXCEEDED:
            self._record_exhaustion_event(cost_metadata, "daily", provider, model)

        return worst

    def _evaluate_budget(
        self,
        *,
        used: float,
        limit: float | None,
        budget_type: str,
    ) -> BudgetCheckResult:
        """Evaluate budget status for a given usage and limit."""
        if limit is None:
            return BudgetCheckResult(
                status=BudgetStatus.UNDER_BUDGET,
                used=used,
                limit=None,
                remaining=None,
                budget_type=budget_type,
            )

        if limit <= 0.0:
            return BudgetCheckResult(
                status=BudgetStatus.EXCEEDED,
                used=used,
                limit=limit,
                remaining=0.0,
                budget_type=budget_type,
            )

        ratio = used / limit
        remaining = max(0.0, limit - used)

        if ratio >= self.BLOCK_THRESHOLD:
            status = BudgetStatus.EXCEEDED
        elif ratio >= self.WARN_THRESHOLD:
            status = BudgetStatus.WARNING
        else:
            status = BudgetStatus.UNDER_BUDGET

        return BudgetCheckResult(
            status=status,
            used=used,
            limit=limit,
            remaining=remaining,
            budget_type=budget_type,
        )

    def _record_exhaustion_event(
        self,
        cost_metadata: CostMetadata,
        budget_type: str,
        provider: str,
        model: str,
    ) -> None:
        """Persist a budget exhaustion event to cost metadata."""
        limit = self._budget_per_plan if budget_type == "plan" else self._budget_per_day
        if limit is None:
            return

        event = BudgetExhaustionEvent(
            timestamp=datetime.now(),
            budget_type=budget_type,
            limit=limit,
            used=cost_metadata.total_cost,
            provider=provider,
            model=model,
        )
        cost_metadata.budget_exhaustion_events.append(event)
        logger.warning(
            "Budget exhausted: %s budget %.4f USD exceeded "
            "(used: %.4f, provider: %s, model: %s)",
            budget_type,
            limit,
            cost_metadata.total_cost,
            provider,
            model,
        )
