"""Cost tracking and budget enforcement for LLM providers.

Tracks token usage and cost per provider/model. Supports soft limits
(warnings) and hard limits (request rejection). Includes BudgetManager
for centralized $500/mo default budget enforcement. Stdlib-only.
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass
from pathlib import Path

# Per-1K-token pricing (input, output) in USD.
# These are approximate list prices; users can override via set_pricing().
DEFAULT_PRICING: dict[str, dict[str, tuple[float, float]]] = {
    "anthropic": {
        "claude-3-opus": (0.015, 0.075),
        "claude-3-sonnet": (0.003, 0.015),
        "claude-3-haiku": (0.00025, 0.00125),
        "default": (0.003, 0.015),
    },
    "openai": {
        "gpt-4": (0.03, 0.06),
        "gpt-4-turbo": (0.01, 0.03),
        "gpt-3.5-turbo": (0.0005, 0.0015),
        "default": (0.01, 0.03),
    },
    "vllm": {
        "default": (0.0, 0.0),  # Local — no API cost.
    },
}


@dataclass
class UsageRecord:
    """Single usage record for a generation call."""

    timestamp: float
    provider: str
    model: str
    input_tokens: int
    output_tokens: int
    cost_usd: float
    latency_ms: float


@dataclass
class BudgetConfig:
    """Budget configuration with soft and hard limits."""

    hard_limit_usd: float = 0.0  # 0 = unlimited
    soft_limit_usd: float = 0.0  # 0 = no warning
    period_seconds: float = 3600.0  # Budget window (default: 1 hour)


class CostTracker:
    """Track LLM usage costs and enforce budgets.

    Maintains a rolling window of usage records for budget enforcement.
    Provides aggregated stats by provider, model, and time period.
    """

    def __init__(self, budget: BudgetConfig | None = None) -> None:
        self._records: list[UsageRecord] = []
        self._pricing: dict[str, dict[str, tuple[float, float]]] = {
            k: dict(v) for k, v in DEFAULT_PRICING.items()
        }
        self._budget = budget or BudgetConfig()

    def set_pricing(
        self, provider: str, model: str, input_per_1k: float, output_per_1k: float
    ) -> None:
        """Override pricing for a specific provider/model."""
        if provider not in self._pricing:
            self._pricing[provider] = {}
        self._pricing[provider][model] = (input_per_1k, output_per_1k)

    def get_pricing(self, provider: str, model: str) -> tuple[float, float]:
        """Get (input_per_1k, output_per_1k) for a provider/model."""
        provider_pricing = self._pricing.get(provider, {})
        if model in provider_pricing:
            return provider_pricing[model]
        return provider_pricing.get("default", (0.0, 0.0))

    def compute_cost(
        self, provider: str, model: str, input_tokens: int, output_tokens: int
    ) -> float:
        """Compute cost in USD for a generation call."""
        input_rate, output_rate = self.get_pricing(provider, model)
        return (input_tokens * input_rate + output_tokens * output_rate) / 1000.0

    def record(
        self,
        provider: str,
        model: str,
        input_tokens: int,
        output_tokens: int,
        latency_ms: float = 0.0,
    ) -> UsageRecord:
        """Record a usage event and return the record."""
        cost = self.compute_cost(provider, model, input_tokens, output_tokens)
        rec = UsageRecord(
            timestamp=time.time(),
            provider=provider,
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cost_usd=cost,
            latency_ms=latency_ms,
        )
        self._records.append(rec)
        return rec

    def check_budget(self) -> dict:
        """Check current spend against budget limits.

        Returns dict with:
          - spent_usd: total spend in current window
          - hard_limit_usd: configured hard limit (0 = unlimited)
          - soft_limit_usd: configured soft limit (0 = no warning)
          - within_budget: True if under hard limit (or unlimited)
          - warning: True if over soft limit but under hard limit
        """
        window_start = time.time() - self._budget.period_seconds
        window_records = [r for r in self._records if r.timestamp >= window_start]
        spent = sum(r.cost_usd for r in window_records)

        hard = self._budget.hard_limit_usd
        soft = self._budget.soft_limit_usd
        within_budget = hard <= 0 or spent < hard
        warning = soft > 0 and spent >= soft and within_budget

        return {
            "spent_usd": spent,
            "hard_limit_usd": hard,
            "soft_limit_usd": soft,
            "within_budget": within_budget,
            "warning": warning,
            "records_in_window": len(window_records),
        }

    def is_within_budget(self) -> bool:
        """Quick check: are we within the hard budget limit?"""
        if self._budget.hard_limit_usd <= 0:
            return True
        window_start = time.time() - self._budget.period_seconds
        spent = sum(r.cost_usd for r in self._records if r.timestamp >= window_start)
        return spent < self._budget.hard_limit_usd

    def summary(self) -> dict:
        """Aggregate usage summary across all records."""
        if not self._records:
            return {
                "total_cost_usd": 0.0,
                "total_input_tokens": 0,
                "total_output_tokens": 0,
                "total_calls": 0,
                "by_provider": {},
            }

        by_provider: dict[str, dict] = {}
        total_cost = 0.0
        total_input = 0
        total_output = 0

        for rec in self._records:
            total_cost += rec.cost_usd
            total_input += rec.input_tokens
            total_output += rec.output_tokens

            if rec.provider not in by_provider:
                by_provider[rec.provider] = {
                    "cost_usd": 0.0,
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "calls": 0,
                    "models": {},
                }
            bp = by_provider[rec.provider]
            bp["cost_usd"] += rec.cost_usd
            bp["input_tokens"] += rec.input_tokens
            bp["output_tokens"] += rec.output_tokens
            bp["calls"] += 1

            if rec.model not in bp["models"]:
                bp["models"][rec.model] = {"cost_usd": 0.0, "calls": 0}
            bp["models"][rec.model]["cost_usd"] += rec.cost_usd
            bp["models"][rec.model]["calls"] += 1

        return {
            "total_cost_usd": total_cost,
            "total_input_tokens": total_input,
            "total_output_tokens": total_output,
            "total_calls": len(self._records),
            "by_provider": by_provider,
        }

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

    @property
    def records(self) -> list[UsageRecord]:
        """Return all usage records (read-only view)."""
        return list(self._records)


class BudgetManager:
    """Centralized budget management with persistent config.

    Default budget is $500/month. Wraps a CostTracker with monthly
    hard/soft limits and persists config to disk. Stdlib-only.
    """

    DEFAULT_MONTHLY_BUDGET_USD: float = 500.0
    SECONDS_PER_MONTH: float = 30.0 * 24 * 3600  # ~30 days

    def __init__(
        self,
        config_path: str | Path | None = None,
        monthly_budget_usd: float | None = None,
    ) -> None:
        self._config_path = Path(config_path) if config_path else None
        self._monthly_budget = monthly_budget_usd or self.DEFAULT_MONTHLY_BUDGET_USD
        self._tracker = CostTracker(
            budget=BudgetConfig(
                hard_limit_usd=self._monthly_budget,
                soft_limit_usd=self._monthly_budget * 0.8,
                period_seconds=self.SECONDS_PER_MONTH,
            )
        )
        if self._config_path and self._config_path.exists():
            self._load_config()

    @property
    def tracker(self) -> CostTracker:
        """The managed CostTracker instance."""
        return self._tracker

    @property
    def monthly_budget_usd(self) -> float:
        """Current monthly budget in USD."""
        return self._monthly_budget

    def set_budget(self, monthly_usd: float) -> None:
        """Update the monthly budget and reconfigure the tracker."""
        self._monthly_budget = monthly_usd
        self._tracker._budget = BudgetConfig(
            hard_limit_usd=monthly_usd,
            soft_limit_usd=monthly_usd * 0.8,
            period_seconds=self.SECONDS_PER_MONTH,
        )
        self._save_config()

    def record(
        self,
        provider: str,
        model: str,
        input_tokens: int,
        output_tokens: int,
        latency_ms: float = 0.0,
    ) -> UsageRecord:
        """Record usage through the managed tracker."""
        return self._tracker.record(provider, model, input_tokens, output_tokens, latency_ms)

    def status(self) -> dict:
        """Return current budget status with spend and limits."""
        budget_check = self._tracker.check_budget()
        summary = self._tracker.summary()
        daily_budget = self._monthly_budget / 30.0
        return {
            "monthly_budget_usd": self._monthly_budget,
            "spent_usd": round(budget_check["spent_usd"], 6),
            "remaining_usd": round(max(0.0, self._monthly_budget - budget_check["spent_usd"]), 6),
            "within_budget": budget_check["within_budget"],
            "warning": budget_check["warning"],
            "daily_budget_usd": round(daily_budget, 2),
            "total_calls": summary["total_calls"],
            "by_provider": summary["by_provider"],
        }

    def _save_config(self) -> None:
        """Persist budget config to disk."""
        if self._config_path:
            self._config_path.parent.mkdir(parents=True, exist_ok=True)
            data = {"monthly_budget_usd": self._monthly_budget}
            self._config_path.write_text(json.dumps(data, indent=2))

    def _load_config(self) -> None:
        """Load budget config from disk."""
        if self._config_path and self._config_path.exists():
            data = json.loads(self._config_path.read_text())
            budget_usd = data.get("monthly_budget_usd", self.DEFAULT_MONTHLY_BUDGET_USD)
            self._monthly_budget = budget_usd
            self._tracker._budget = BudgetConfig(
                hard_limit_usd=budget_usd,
                soft_limit_usd=budget_usd * 0.8,
                period_seconds=self.SECONDS_PER_MONTH,
            )


class CostAwareRouter:
    """Extends ProviderRouter with cost-aware selection.

    Factors in pricing when selecting providers. For equivalent task
    suitability, prefers the cheaper option. Respects budget limits
    by skipping providers that would exceed the hard cap.
    """

    def __init__(
        self,
        tracker: CostTracker,
        provider_costs: dict[str, float] | None = None,
    ) -> None:
        self._tracker = tracker
        # Estimated cost per 1K tokens (combined) by provider name.
        self._provider_costs: dict[str, float] = provider_costs or {
            "vllm": 0.0,
            "anthropic": 0.018,
            "openai": 0.04,
        }

    def rank_by_cost(self, candidates: list[str]) -> list[str]:
        """Re-rank provider names by estimated cost (cheapest first)."""
        return sorted(candidates, key=lambda n: self._provider_costs.get(n, 999.0))

    def should_skip(self, provider_name: str, estimated_tokens: int = 1000) -> bool:
        """Check if using this provider would exceed the budget."""
        if not self._tracker._budget.hard_limit_usd:
            return False
        cost_per_1k = self._provider_costs.get(provider_name, 0.0)
        estimated_cost = (estimated_tokens * cost_per_1k) / 1000.0
        budget = self._tracker.check_budget()
        return bool(budget["spent_usd"] + estimated_cost >= budget["hard_limit_usd"])

    def select(
        self,
        candidates: list[str],
        task_type: str | None = None,
        estimated_tokens: int = 1000,
    ) -> list[str]:
        """Select and rank candidates considering cost + budget.

        Returns ordered list of provider names (best first).
        Filters out providers that would exceed budget.
        """
        # Start with cost ranking.
        ranked = self.rank_by_cost(candidates)

        # Filter out budget-busting providers.
        affordable = [name for name in ranked if not self.should_skip(name, estimated_tokens)]

        return affordable if affordable else ranked[:1]
