"""Cost tracking -- per-request logging, savings calculation, and cloud pricing data.

Uses aiosqlite for async SQLite storage.  The CostTracker is initialised once
at startup and shared across the application lifetime.
"""

from __future__ import annotations

import logging
import uuid
from datetime import datetime, timezone
from typing import TYPE_CHECKING

import aiosqlite
from pydantic import BaseModel

if TYPE_CHECKING:
    from pathlib import Path

    from llmhosts.router.models import RoutingDecision

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Cloud pricing table (USD per token, updated 2025/2026 rates)
# ---------------------------------------------------------------------------

CLOUD_PRICING: dict[str, dict[str, float]] = {
    # OpenAI
    "gpt-4o": {"input": 2.50 / 1e6, "output": 10.00 / 1e6},
    "gpt-4o-mini": {"input": 0.15 / 1e6, "output": 0.60 / 1e6},
    "gpt-4-turbo": {"input": 10.00 / 1e6, "output": 30.00 / 1e6},
    "gpt-4-turbo-2024-04-09": {"input": 10.00 / 1e6, "output": 30.00 / 1e6},
    "gpt-3.5-turbo": {"input": 0.50 / 1e6, "output": 1.50 / 1e6},
    "gpt-3.5-turbo-0125": {"input": 0.50 / 1e6, "output": 1.50 / 1e6},
    "o1": {"input": 15.00 / 1e6, "output": 60.00 / 1e6},
    "o1-mini": {"input": 3.00 / 1e6, "output": 12.00 / 1e6},
    "o3-mini": {"input": 1.10 / 1e6, "output": 4.40 / 1e6},
    # Anthropic
    "claude-3-5-sonnet-20241022": {"input": 3.00 / 1e6, "output": 15.00 / 1e6},
    "claude-3-5-haiku-20241022": {"input": 0.80 / 1e6, "output": 4.00 / 1e6},
    "claude-3-opus-20240229": {"input": 15.00 / 1e6, "output": 75.00 / 1e6},
    "claude-3-sonnet-20240229": {"input": 3.00 / 1e6, "output": 15.00 / 1e6},
    "claude-3-haiku-20240307": {"input": 0.25 / 1e6, "output": 1.25 / 1e6},
    # Mistral (via API)
    "mistral-large-latest": {"input": 2.00 / 1e6, "output": 6.00 / 1e6},
    "mistral-small-latest": {"input": 0.20 / 1e6, "output": 0.60 / 1e6},
    # Google
    "gemini-1.5-pro": {"input": 1.25 / 1e6, "output": 5.00 / 1e6},
    "gemini-1.5-flash": {"input": 0.075 / 1e6, "output": 0.30 / 1e6},
}

# Default "reference" model for cloud-equivalent cost calculation when exact model
# pricing is unknown. We use gpt-4o-mini as a conservative baseline.
_DEFAULT_REFERENCE_MODEL = "gpt-4o-mini"


# ---------------------------------------------------------------------------
# Report models
# ---------------------------------------------------------------------------


class SavingsReport(BaseModel):
    """Aggregate savings for a time period."""

    period: str
    total_requests: int = 0
    local_requests: int = 0
    cloud_requests: int = 0
    cache_hits: int = 0
    actual_cost: float = 0.0
    cloud_equivalent_cost: float = 0.0
    savings: float = 0.0
    savings_percent: float = 0.0


class RequestLog(BaseModel):
    """A single logged request with routing and cost data."""

    request_id: str
    timestamp: datetime
    model_requested: str
    model_used: str
    backend_type: str
    routing_tier: str
    routing_reasoning: str
    input_tokens: int
    output_tokens: int
    actual_cost: float
    cloud_equivalent_cost: float
    latency_ms: float
    cached: bool


# ---------------------------------------------------------------------------
# Helper functions
# ---------------------------------------------------------------------------


def estimate_tokens(text: str) -> int:
    """Rough token count: ~4 characters per token on average."""
    return max(1, len(text) // 4)


def estimate_request_tokens(messages: list[dict[str, object]]) -> int:
    """Estimate total input tokens from a list of message dicts."""
    total_chars = 0
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            total_chars += len(content)
        elif isinstance(content, list):
            for block in content:
                if isinstance(block, dict):
                    total_chars += len(block.get("text", ""))
    return max(1, total_chars // 4)


def cloud_equivalent_cost(model_requested: str, input_tokens: int, output_tokens: int) -> float:
    """Calculate what this request would have cost on the cloud.

    Looks up the model in :data:`CLOUD_PRICING`.  Falls back to the default
    reference model if the exact model is unknown.
    """
    pricing = CLOUD_PRICING.get(model_requested)
    if pricing is None:
        pricing = CLOUD_PRICING.get(_DEFAULT_REFERENCE_MODEL, {"input": 0.0, "output": 0.0})
    return input_tokens * pricing["input"] + output_tokens * pricing["output"]


def actual_request_cost(backend_type: str, model: str, input_tokens: int, output_tokens: int) -> float:
    """Calculate actual cost for a request.

    Local backends (ollama) are free.  Cloud backends use :data:`CLOUD_PRICING`.
    """
    if backend_type == "ollama":
        return 0.0
    pricing = CLOUD_PRICING.get(model)
    if pricing is None:
        return 0.0  # Unknown model, can't calculate
    return input_tokens * pricing["input"] + output_tokens * pricing["output"]


# ---------------------------------------------------------------------------
# SQL schema
# ---------------------------------------------------------------------------

_CREATE_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS request_log (
    request_id     TEXT PRIMARY KEY,
    timestamp      TEXT NOT NULL,
    model_requested TEXT NOT NULL,
    model_used     TEXT NOT NULL,
    backend_type   TEXT NOT NULL,
    routing_tier   TEXT NOT NULL,
    routing_reasoning TEXT NOT NULL DEFAULT '',
    input_tokens   INTEGER NOT NULL DEFAULT 0,
    output_tokens  INTEGER NOT NULL DEFAULT 0,
    actual_cost    REAL NOT NULL DEFAULT 0.0,
    cloud_equivalent_cost REAL NOT NULL DEFAULT 0.0,
    latency_ms     REAL NOT NULL DEFAULT 0.0,
    cached         INTEGER NOT NULL DEFAULT 0
);
"""

_CREATE_INDEX_SQL = """
CREATE INDEX IF NOT EXISTS idx_request_log_timestamp ON request_log (timestamp);
CREATE INDEX IF NOT EXISTS idx_request_log_backend ON request_log (backend_type);
"""


# ---------------------------------------------------------------------------
# CostTracker
# ---------------------------------------------------------------------------


class CostTracker:
    """Tracks per-request and aggregate costs via async SQLite.

    Usage::

        tracker = CostTracker(db_path=Path("~/.llmhosts/cost.db"))
        await tracker.initialize()
        await tracker.record_request(...)
        report = await tracker.get_savings("month")
        await tracker.close()
    """

    def __init__(self, db_path: Path) -> None:
        self._db_path = db_path
        self._db: aiosqlite.Connection | None = None

    async def initialize(self) -> None:
        """Open the database and create tables if they don't exist."""
        self._db_path.parent.mkdir(parents=True, exist_ok=True)
        self._db = await aiosqlite.connect(str(self._db_path))
        await self._db.execute("PRAGMA journal_mode=WAL")
        await self._db.execute("PRAGMA synchronous=NORMAL")
        await self._db.executescript(_CREATE_TABLE_SQL + _CREATE_INDEX_SQL)
        await self._db.commit()
        logger.info("CostTracker initialised: %s", self._db_path)

    async def close(self) -> None:
        """Close the database connection."""
        if self._db is not None:
            await self._db.close()
            self._db = None

    async def record_request(
        self,
        decision: RoutingDecision,
        model_requested: str,
        input_tokens: int,
        output_tokens: int,
        actual_cost: float,
        latency_ms: float,
        cached: bool = False,
        request_id: str | None = None,
    ) -> str:
        """Record a completed request with cost data.

        Returns the generated request_id.
        """
        assert self._db is not None, "CostTracker not initialised -- call initialize() first"
        rid = request_id or uuid.uuid4().hex[:16]
        now = datetime.now(tz=timezone.utc).isoformat()
        equiv = cloud_equivalent_cost(model_requested, input_tokens, output_tokens)

        await self._db.execute(
            """
            INSERT OR REPLACE INTO request_log
                (request_id, timestamp, model_requested, model_used, backend_type,
                 routing_tier, routing_reasoning, input_tokens, output_tokens,
                 actual_cost, cloud_equivalent_cost, latency_ms, cached)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                rid,
                now,
                model_requested,
                decision.model,
                decision.backend_type,
                decision.tier,
                decision.reasoning,
                input_tokens,
                output_tokens,
                actual_cost,
                equiv,
                latency_ms,
                1 if cached else 0,
            ),
        )
        await self._db.commit()
        return rid

    async def get_savings(self, period: str = "all") -> SavingsReport:
        """Calculate savings for a period.

        Parameters
        ----------
        period:
            One of ``"day"``, ``"week"``, ``"month"``, ``"all"``.

        Returns
        -------
        SavingsReport
        """
        assert self._db is not None, "CostTracker not initialised"
        where_clause, params = self._period_filter(period)

        query = f"""
            SELECT
                COUNT(*) AS total_requests,
                SUM(CASE WHEN backend_type = 'ollama' THEN 1 ELSE 0 END) AS local_requests,
                SUM(CASE WHEN backend_type != 'ollama' THEN 1 ELSE 0 END) AS cloud_requests,
                SUM(CASE WHEN cached = 1 THEN 1 ELSE 0 END) AS cache_hits,
                COALESCE(SUM(actual_cost), 0.0) AS actual_cost,
                COALESCE(SUM(cloud_equivalent_cost), 0.0) AS cloud_equivalent_cost
            FROM request_log
            {where_clause}
        """  # nosec B608 -- where_clause is built from validated period enum, values use ? params

        async with self._db.execute(query, params) as cursor:
            row = await cursor.fetchone()

        if row is None or row[0] == 0:
            return SavingsReport(period=period)

        total_requests, local_requests, cloud_requests, cache_hits, act_cost, equiv_cost = row
        savings = equiv_cost - act_cost
        savings_pct = (savings / equiv_cost * 100.0) if equiv_cost > 0 else 0.0

        return SavingsReport(
            period=period,
            total_requests=total_requests,
            local_requests=local_requests,
            cloud_requests=cloud_requests,
            cache_hits=cache_hits,
            actual_cost=round(act_cost, 6),
            cloud_equivalent_cost=round(equiv_cost, 6),
            savings=round(savings, 6),
            savings_percent=round(savings_pct, 2),
        )

    async def get_request_log(self, limit: int = 100, offset: int = 0) -> list[RequestLog]:
        """Get recent requests with routing decisions and costs."""
        assert self._db is not None, "CostTracker not initialised"

        query = """
            SELECT request_id, timestamp, model_requested, model_used, backend_type,
                   routing_tier, routing_reasoning, input_tokens, output_tokens,
                   actual_cost, cloud_equivalent_cost, latency_ms, cached
            FROM request_log
            ORDER BY timestamp DESC
            LIMIT ? OFFSET ?
        """

        results: list[RequestLog] = []
        async with self._db.execute(query, (limit, offset)) as cursor:
            async for row in cursor:
                results.append(
                    RequestLog(
                        request_id=row[0],
                        timestamp=datetime.fromisoformat(row[1]),
                        model_requested=row[2],
                        model_used=row[3],
                        backend_type=row[4],
                        routing_tier=row[5],
                        routing_reasoning=row[6],
                        input_tokens=row[7],
                        output_tokens=row[8],
                        actual_cost=row[9],
                        cloud_equivalent_cost=row[10],
                        latency_ms=row[11],
                        cached=bool(row[12]),
                    )
                )
        return results

    async def get_total_savings(self) -> float:
        """Quick helper: total savings across all time."""
        report = await self.get_savings("all")
        return report.savings

    # -- internal helpers -----------------------------------------------------

    @staticmethod
    def _period_filter(period: str) -> tuple[str, tuple[str, ...]]:
        """Build a SQL WHERE clause for the given period.

        Returns ``(where_clause_str, params_tuple)``.
        """
        if period == "all":
            return ("", ())

        # Map period name to SQLite date modifier
        modifiers = {
            "day": "-1 day",
            "week": "-7 days",
            "month": "-30 days",
        }
        modifier = modifiers.get(period)
        if modifier is None:
            logger.warning("Unknown period '%s', returning all-time data", period)
            return ("", ())

        return ("WHERE timestamp >= datetime('now', ?)", (modifier,))
