"""Real-time LLM cost intelligence — compares local/cached vs cloud API pricing."""

from __future__ import annotations

import dataclasses
import datetime
import logging
from typing import ClassVar

logger = logging.getLogger(__name__)

# Public cloud pricing per million tokens (input/output), USD. Updated Feb 2026.
CLOUD_PRICES: dict[str, dict[str, float]] = {
    "gpt-4o": {"input": 2.50, "output": 10.00},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "claude-3-5-sonnet": {"input": 3.00, "output": 15.00},
    "claude-3-haiku": {"input": 0.25, "output": 1.25},
    "gemini-1.5-pro": {"input": 1.25, "output": 5.00},
    "gemini-1.5-flash": {"input": 0.075, "output": 0.30},
    "llama-3.3-70b": {"input": 0.59, "output": 0.79},  # Groq pricing
    "mistral-large": {"input": 2.00, "output": 6.00},
}

# Model family → cloud baseline for Ollama-style names (e.g. "llama3.3:70b")
_FAMILY_PRICES: dict[str, dict[str, float]] = {
    "llama": {"input": 0.59, "output": 0.79},  # Groq Llama pricing
    "mistral": {"input": 0.50, "output": 1.50},  # Mistral 7B API
    "gemma": {"input": 0.15, "output": 0.60},  # GPT-4o-mini tier
    "phi": {"input": 0.15, "output": 0.60},
    "qwen": {"input": 0.15, "output": 0.60},
    "deepseek": {"input": 0.27, "output": 1.10},
}

# Electricity cost per kWh (USD). Configurable.
DEFAULT_ELECTRICITY_KWH = 0.12


@dataclasses.dataclass
class SavingsReport:
    model: str
    input_tokens: int
    output_tokens: int
    cloud_cost_usd: float
    local_cost_usd: float  # electricity only
    savings_usd: float
    savings_pct: float
    timestamp: datetime.datetime = dataclasses.field(
        default_factory=lambda: datetime.datetime.now(datetime.timezone.utc)
    )

    def to_dict(self) -> dict:
        return {
            "model": self.model,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "cloud_cost_usd": round(self.cloud_cost_usd, 6),
            "local_cost_usd": round(self.local_cost_usd, 6),
            "savings_usd": round(self.savings_usd, 6),
            "savings_pct": round(self.savings_pct, 2),
            "timestamp": self.timestamp.isoformat(),
        }


class CostOracle:
    """Calculates real-time savings of local inference vs cloud APIs."""

    WATT_HOUR_PER_TOKEN: ClassVar[dict[str, float]] = {
        # Watt-hours per 1000 tokens (inference only)
        "7b": 0.0003,
        "13b": 0.0006,
        "34b": 0.0015,
        "70b": 0.0030,
        "default": 0.0003,
    }

    def __init__(self, electricity_kwh: float = DEFAULT_ELECTRICITY_KWH) -> None:
        self._electricity_kwh = electricity_kwh
        self._reports: list[SavingsReport] = []

    def _electricity_cost(self, model: str, total_tokens: int) -> float:
        """Estimate electricity cost for local inference."""
        size_key = "default"
        for key in self.WATT_HOUR_PER_TOKEN:
            if key in model.lower():
                size_key = key
                break
        wh_per_1k = self.WATT_HOUR_PER_TOKEN[size_key]
        wh = (total_tokens / 1000) * wh_per_1k
        kwh = wh / 1000
        return kwh * self._electricity_kwh

    def _cloud_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        """Look up cloud API cost. Falls back to cheapest known price."""
        # Strip provider prefix (e.g. "ollama/", "local/") and tag (e.g. ":latest", ":70b")
        base = model.lower().split("/")[-1].split(":")[0]
        # Normalize separators: "llama3.3" → "llama3-3", then try exact match
        norm = base.replace(".", "-")

        prices = CLOUD_PRICES.get(norm)
        if not prices:
            # Try exact match on base (without separator normalisation)
            prices = CLOUD_PRICES.get(base)
        if not prices:
            # Substring match against CLOUD_PRICES keys
            for key, p in CLOUD_PRICES.items():
                if key in norm or norm in key:
                    prices = p
                    break
        if not prices:
            # Family match for Ollama-style names (e.g. "llama3.3" → "llama" family)
            for family, p in _FAMILY_PRICES.items():
                if base.startswith(family):
                    prices = p
                    logger.debug("Cost oracle: matched '%s' to family '%s'", model, family)
                    break
        if not prices:
            # Final fallback: GPT-4o-mini as conservative baseline
            prices = CLOUD_PRICES["gpt-4o-mini"]
            logger.debug("Cost oracle: no match for '%s', using gpt-4o-mini baseline", model)

        return (input_tokens * prices["input"] + output_tokens * prices["output"]) / 1_000_000

    def calculate(self, model: str, input_tokens: int, output_tokens: int) -> SavingsReport:
        """Calculate savings for a single inference call."""
        total = input_tokens + output_tokens
        cloud = self._cloud_cost(model, input_tokens, output_tokens)
        local = self._electricity_cost(model, total)
        savings = max(0.0, cloud - local)
        pct = (savings / cloud * 100) if cloud > 0 else 0.0
        report = SavingsReport(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cloud_cost_usd=cloud,
            local_cost_usd=local,
            savings_usd=savings,
            savings_pct=pct,
        )
        self._reports.append(report)
        return report

    def cumulative_savings(self) -> float:
        """Total savings across all tracked calls."""
        return sum(r.savings_usd for r in self._reports)

    def session_summary(self) -> dict:
        """Return a summary dict for the current session."""
        if not self._reports:
            return {"total_calls": 0, "total_savings_usd": 0.0, "avg_savings_pct": 0.0, "reports": []}
        return {
            "total_calls": len(self._reports),
            "total_savings_usd": round(self.cumulative_savings(), 4),
            "avg_savings_pct": round(sum(r.savings_pct for r in self._reports) / len(self._reports), 2),
            "reports": [r.to_dict() for r in self._reports],
        }
