"""Cost estimator — predict LLM call costs before execution."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any

logger = logging.getLogger(__name__)

# Pricing per 1M tokens (USD) — input / output
# Updated manually; can be overridden via config.
_DEFAULT_PRICING: dict[str, tuple[float, float]] = {
    # Anthropic
    "anthropic/claude-opus-4-20250514": (15.0, 75.0),
    "anthropic/claude-sonnet-4-20250514": (3.0, 15.0),
    "anthropic/claude-haiku-3": (0.25, 1.25),
    # OpenAI
    "openai/gpt-4o": (2.5, 10.0),
    "openai/gpt-4o-mini": (0.15, 0.60),
    "openai/gpt-4-turbo": (10.0, 30.0),
    # Google
    "google/gemini-pro": (0.50, 1.50),
    "google/gemini-1.5-pro": (3.50, 10.50),
    "google/gemini-1.5-flash": (0.075, 0.30),
    # Groq (free tier, but track anyway)
    "groq/llama-3-8b": (0.05, 0.08),
    "groq/llama-3-70b": (0.59, 0.79),
    "groq/mixtral-8x7b": (0.24, 0.24),
    # Mistral
    "mistral/mistral-large": (3.0, 9.0),
    "mistral/mistral-small": (1.0, 3.0),
}

# Rough estimates for tokens per character
_CHARS_PER_TOKEN = 4


@dataclass
class CostEstimate:
    """Pre-execution cost estimate for an LLM call."""

    model: str
    estimated_input_tokens: int
    estimated_output_tokens: int
    estimated_cost_usd: float
    price_per_m_input: float
    price_per_m_output: float

    @property
    def formatted(self) -> str:
        return (
            f"~{self.estimated_input_tokens:,} in + "
            f"~{self.estimated_output_tokens:,} out → "
            f"${self.estimated_cost_usd:.4f}"
        )


@dataclass
class CostComparison:
    """Compare estimated vs actual cost after a call."""

    model: str
    estimated_cost: float
    actual_cost: float

    @property
    def difference(self) -> float:
        return self.actual_cost - self.estimated_cost

    @property
    def accuracy_pct(self) -> float:
        if self.estimated_cost == 0:
            return 0.0
        return (1 - abs(self.difference) / self.estimated_cost) * 100


class CostEstimator:
    """Estimate LLM call costs before execution.

    Uses model-specific pricing tables and message size heuristics
    to predict costs. Supports custom pricing overrides.
    """

    def __init__(
        self,
        pricing_overrides: dict[str, tuple[float, float]] | None = None,
    ) -> None:
        self._pricing = dict(_DEFAULT_PRICING)
        if pricing_overrides:
            self._pricing.update(pricing_overrides)

    def estimate(
        self,
        model: str,
        messages: list[dict[str, Any]] | None = None,
        input_tokens: int | None = None,
        expected_output_tokens: int = 500,
    ) -> CostEstimate:
        """Estimate the cost of an LLM call.

        Provide either messages (will estimate tokens from content length)
        or input_tokens directly.
        """
        price_in, price_out = self._get_pricing(model)

        if input_tokens is None:
            input_tokens = self._estimate_message_tokens(messages or [])

        cost = (input_tokens / 1_000_000) * price_in + (
            expected_output_tokens / 1_000_000
        ) * price_out

        return CostEstimate(
            model=model,
            estimated_input_tokens=input_tokens,
            estimated_output_tokens=expected_output_tokens,
            estimated_cost_usd=cost,
            price_per_m_input=price_in,
            price_per_m_output=price_out,
        )

    def compare(
        self,
        estimate: CostEstimate,
        actual_input_tokens: int,
        actual_output_tokens: int,
    ) -> CostComparison:
        """Compare an estimate against actual usage."""
        actual_cost = (actual_input_tokens / 1_000_000) * estimate.price_per_m_input + (
            actual_output_tokens / 1_000_000
        ) * estimate.price_per_m_output
        return CostComparison(
            model=estimate.model,
            estimated_cost=estimate.estimated_cost_usd,
            actual_cost=actual_cost,
        )

    def calculate_cost(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
    ) -> float:
        """Calculate the exact cost for a completed call."""
        price_in, price_out = self._get_pricing(model)
        return (input_tokens / 1_000_000) * price_in + (output_tokens / 1_000_000) * price_out

    def get_pricing(self, model: str) -> tuple[float, float]:
        """Get pricing for a model. Returns (input_per_M, output_per_M)."""
        return self._get_pricing(model)

    def update_pricing(self, model: str, input_price: float, output_price: float) -> None:
        """Update or add pricing for a model."""
        self._pricing[model] = (input_price, output_price)

    def _get_pricing(self, model: str) -> tuple[float, float]:
        """Look up pricing, with fallback for unknown models."""
        if model in self._pricing:
            return self._pricing[model]

        # Try matching by prefix (e.g., "anthropic/claude-sonnet"
        # matches "anthropic/claude-sonnet-4-...")
        for key, prices in self._pricing.items():
            if model.startswith(key.rsplit("-", 1)[0]):
                return prices

        # Unknown model — use a conservative mid-range estimate
        logger.warning("No pricing data for model '%s', using fallback", model)
        return (3.0, 15.0)

    def _estimate_message_tokens(self, messages: list[dict[str, Any]]) -> int:
        """Estimate token count from messages using character count heuristic."""
        total_chars = 0
        for msg in messages:
            content = msg.get("content", "")
            if isinstance(content, str):
                total_chars += len(content)
            elif isinstance(content, list):
                for part in content:
                    if isinstance(part, dict):
                        total_chars += len(part.get("text", ""))

        # Add overhead for message structure (~4 tokens per message)
        overhead = len(messages) * 4
        return (total_chars // _CHARS_PER_TOKEN) + overhead
