"""Cloud API budget tracker for the three-tier model router.

Tracks cloud API spend (daily / monthly) and enforces configurable limits.
SQLite-backed with in-memory fallback — degrades gracefully if the database
is unavailable.

References:
  - Issue #882 — Model Tiering Router: Local 8B / Hermes 70B / Cloud API Cascade
"""

import logging
import sqlite3
import threading
import time
from dataclasses import dataclass
from datetime import UTC, date, datetime
from pathlib import Path

from config import settings

logger = logging.getLogger(__name__)

# ── Cost estimates (USD per 1 K tokens, input / output) ──────────────────────
# Updated 2026-03.  Estimates only — actual costs vary by tier/usage.
_COST_PER_1K: dict[str, dict[str, float]] = {
    # Claude models
    "claude-haiku-4-5": {"input": 0.00025, "output": 0.00125},
    "claude-sonnet-4-5": {"input": 0.003, "output": 0.015},
    "claude-opus-4-5": {"input": 0.015, "output": 0.075},
    "haiku": {"input": 0.00025, "output": 0.00125},
    "sonnet": {"input": 0.003, "output": 0.015},
    "opus": {"input": 0.015, "output": 0.075},
    # GPT-4o
    "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
    "gpt-4o": {"input": 0.0025, "output": 0.01},
    # Grok (xAI)
    "grok-3-fast": {"input": 0.003, "output": 0.015},
    "grok-3": {"input": 0.005, "output": 0.025},
}
_DEFAULT_COST: dict[str, float] = {"input": 0.003, "output": 0.015}  # conservative fallback


def estimate_cost_usd(model: str, tokens_in: int, tokens_out: int) -> float:
    """Estimate the cost of a single request in USD.

    Matches the model name by substring so versioned names like
    ``claude-haiku-4-5-20251001`` still resolve correctly.

    Args:
        model:      Model name as passed to the provider.
        tokens_in:  Number of input (prompt) tokens consumed.
        tokens_out: Number of output (completion) tokens generated.

    Returns:
        Estimated cost in USD (may be zero for unknown models).
    """
    model_lower = model.lower()
    rates = _DEFAULT_COST
    for key, rate in _COST_PER_1K.items():
        if key in model_lower:
            rates = rate
            break
    return (tokens_in * rates["input"] + tokens_out * rates["output"]) / 1000.0


@dataclass
class SpendRecord:
    """A single spend event."""

    ts: float
    provider: str
    model: str
    tokens_in: int
    tokens_out: int
    cost_usd: float
    tier: str


class BudgetTracker:
    """Tracks cloud API spend with configurable daily / monthly limits.

    Persists spend records to SQLite (``data/budget.db`` by default).
    Falls back to in-memory tracking when the database is unavailable —
    budget enforcement still works; records are lost on restart.

    Limits are read from ``settings``:

    * ``tier_cloud_daily_budget_usd``   — daily ceiling (0 = disabled)
    * ``tier_cloud_monthly_budget_usd`` — monthly ceiling (0 = disabled)

    Usage::

        tracker = BudgetTracker()

        if tracker.cloud_allowed():
            # … make cloud API call …
            tracker.record_spend("anthropic", "claude-haiku-4-5", 100, 200)

        summary = tracker.get_summary()
        print(summary["daily_usd"], "/", summary["daily_limit_usd"])
    """

    _DB_PATH = "data/budget.db"

    def __init__(self, db_path: str | None = None) -> None:
        """Initialise the tracker.

        Args:
            db_path: Path to the SQLite database.  Defaults to
                     ``data/budget.db``.  Pass ``":memory:"`` for tests.
        """
        self._db_path = db_path or self._DB_PATH
        self._lock = threading.Lock()
        self._in_memory: list[SpendRecord] = []
        self._db_ok = False
        self._init_db()

    # ── Database initialisation ──────────────────────────────────────────────

    def _init_db(self) -> None:
        """Create the spend table (and parent directory) if needed."""
        try:
            if self._db_path != ":memory:":
                Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
            with self._connect() as conn:
                conn.execute(
                    """
                    CREATE TABLE IF NOT EXISTS cloud_spend (
                        id         INTEGER PRIMARY KEY AUTOINCREMENT,
                        ts         REAL    NOT NULL,
                        provider   TEXT    NOT NULL,
                        model      TEXT    NOT NULL,
                        tokens_in  INTEGER NOT NULL DEFAULT 0,
                        tokens_out INTEGER NOT NULL DEFAULT 0,
                        cost_usd   REAL    NOT NULL DEFAULT 0.0,
                        tier       TEXT    NOT NULL DEFAULT 'cloud'
                    )
                    """
                )
                conn.execute("CREATE INDEX IF NOT EXISTS idx_spend_ts ON cloud_spend(ts)")
            self._db_ok = True
            logger.debug("BudgetTracker: SQLite initialised at %s", self._db_path)
        except Exception as exc:
            logger.warning("BudgetTracker: SQLite unavailable, using in-memory fallback: %s", exc)

    def _connect(self) -> sqlite3.Connection:
        return sqlite3.connect(self._db_path, timeout=5)

    # ── Public API ───────────────────────────────────────────────────────────

    def record_spend(
        self,
        provider: str,
        model: str,
        tokens_in: int = 0,
        tokens_out: int = 0,
        cost_usd: float | None = None,
        tier: str = "cloud",
    ) -> float:
        """Record a cloud API spend event and return the cost recorded.

        Args:
            provider:   Provider name (e.g. ``"anthropic"``, ``"openai"``).
            model:      Model name used for the request.
            tokens_in:  Input token count (prompt).
            tokens_out: Output token count (completion).
            cost_usd:   Explicit cost override.  If ``None``, the cost is
                        estimated from the token counts and model rates.
            tier:       Tier label for the request (default ``"cloud"``).

        Returns:
            The cost recorded in USD.
        """
        if cost_usd is None:
            cost_usd = estimate_cost_usd(model, tokens_in, tokens_out)

        ts = time.time()
        record = SpendRecord(ts, provider, model, tokens_in, tokens_out, cost_usd, tier)

        with self._lock:
            if self._db_ok:
                try:
                    with self._connect() as conn:
                        conn.execute(
                            """
                            INSERT INTO cloud_spend
                                (ts, provider, model, tokens_in, tokens_out, cost_usd, tier)
                            VALUES (?, ?, ?, ?, ?, ?, ?)
                            """,
                            (ts, provider, model, tokens_in, tokens_out, cost_usd, tier),
                        )
                    logger.debug(
                        "BudgetTracker: recorded %.6f USD (%s/%s, in=%d out=%d tier=%s)",
                        cost_usd,
                        provider,
                        model,
                        tokens_in,
                        tokens_out,
                        tier,
                    )
                    return cost_usd
                except Exception as exc:
                    logger.warning("BudgetTracker: DB write failed, falling back: %s", exc)
            self._in_memory.append(record)

        return cost_usd

    def get_daily_spend(self) -> float:
        """Return total cloud spend for the current UTC day in USD."""
        today = date.today()
        since = datetime(today.year, today.month, today.day, tzinfo=UTC).timestamp()
        return self._query_spend(since)

    def get_monthly_spend(self) -> float:
        """Return total cloud spend for the current UTC month in USD."""
        today = date.today()
        since = datetime(today.year, today.month, 1, tzinfo=UTC).timestamp()
        return self._query_spend(since)

    def cloud_allowed(self) -> bool:
        """Return ``True`` if cloud API spend is within configured limits.

        Checks both daily and monthly ceilings.  A limit of ``0`` disables
        that particular check.
        """
        daily_limit = settings.tier_cloud_daily_budget_usd
        monthly_limit = settings.tier_cloud_monthly_budget_usd

        if daily_limit > 0:
            daily_spend = self.get_daily_spend()
            if daily_spend >= daily_limit:
                logger.warning(
                    "BudgetTracker: daily cloud budget exhausted (%.4f / %.4f USD)",
                    daily_spend,
                    daily_limit,
                )
                return False

        if monthly_limit > 0:
            monthly_spend = self.get_monthly_spend()
            if monthly_spend >= monthly_limit:
                logger.warning(
                    "BudgetTracker: monthly cloud budget exhausted (%.4f / %.4f USD)",
                    monthly_spend,
                    monthly_limit,
                )
                return False

        return True

    def get_summary(self) -> dict:
        """Return a spend summary dict suitable for dashboards / logging.

        Keys: ``daily_usd``, ``monthly_usd``, ``daily_limit_usd``,
        ``monthly_limit_usd``, ``daily_ok``, ``monthly_ok``.
        """
        daily = self.get_daily_spend()
        monthly = self.get_monthly_spend()
        daily_limit = settings.tier_cloud_daily_budget_usd
        monthly_limit = settings.tier_cloud_monthly_budget_usd
        return {
            "daily_usd": round(daily, 6),
            "monthly_usd": round(monthly, 6),
            "daily_limit_usd": daily_limit,
            "monthly_limit_usd": monthly_limit,
            "daily_ok": daily_limit <= 0 or daily < daily_limit,
            "monthly_ok": monthly_limit <= 0 or monthly < monthly_limit,
        }

    # ── Internal helpers ─────────────────────────────────────────────────────

    def _query_spend(self, since_ts: float) -> float:
        """Sum ``cost_usd`` for records with ``ts >= since_ts``."""
        if self._db_ok:
            try:
                with self._connect() as conn:
                    row = conn.execute(
                        "SELECT COALESCE(SUM(cost_usd), 0.0) FROM cloud_spend WHERE ts >= ?",
                        (since_ts,),
                    ).fetchone()
                    return float(row[0]) if row else 0.0
            except Exception as exc:
                logger.warning("BudgetTracker: DB read failed: %s", exc)
        # In-memory fallback
        return sum(r.cost_usd for r in self._in_memory if r.ts >= since_ts)


# ── Module-level singleton ────────────────────────────────────────────────────

_budget_tracker: BudgetTracker | None = None


def get_budget_tracker() -> BudgetTracker:
    """Get or create the module-level BudgetTracker singleton."""
    global _budget_tracker
    if _budget_tracker is None:
        _budget_tracker = BudgetTracker()
    return _budget_tracker
