"""Per-provider cost table with default token cost estimates.

Provides offline cost estimation for token usage across supported
providers and models. Costs are in USD per token.

Default estimates are approximations for common models and should be
updated as provider pricing changes.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import ClassVar


@dataclass(frozen=True, slots=True)
class CostEntry:
    """Cost per token for a specific provider/model combination.

    Attributes:
        input_cost_per_token: USD cost per input/prompt token.
        output_cost_per_token: USD cost per output/completion token.
    """

    input_cost_per_token: float
    output_cost_per_token: float

    def __post_init__(self) -> None:
        if self.input_cost_per_token < 0.0:
            raise ValueError(
                f"input_cost_per_token must be >= 0, got {self.input_cost_per_token}"
            )
        if self.output_cost_per_token < 0.0:
            raise ValueError(
                f"output_cost_per_token must be >= 0, got {self.output_cost_per_token}"
            )

    def estimate_cost(self, input_tokens: int, output_tokens: int) -> float:
        """Estimate the cost for a given token count.

        Args:
            input_tokens: Number of input tokens.
            output_tokens: Number of output tokens.

        Returns:
            Estimated cost in USD.

        Raises:
            ValueError: If token counts are negative.
        """
        if input_tokens < 0:
            raise ValueError(f"input_tokens must be >= 0, got {input_tokens}")
        if output_tokens < 0:
            raise ValueError(f"output_tokens must be >= 0, got {output_tokens}")
        return (
            self.input_cost_per_token * input_tokens
            + self.output_cost_per_token * output_tokens
        )


class ProviderCostTable:
    """Lookup table for provider/model cost estimates.

    Maintains a registry of known cost entries and provides a fallback
    default for unknown models. Thread-safe for read operations.

    Example::

        table = ProviderCostTable()
        entry = table.get_cost_entry("openai", "gpt-4o")
        cost = entry.estimate_cost(input_tokens=1000, output_tokens=500)
    """

    # Default fallback cost for unknown models (conservative estimate)
    DEFAULT_COST: ClassVar[CostEntry] = CostEntry(
        input_cost_per_token=0.000005,
        output_cost_per_token=0.000015,
    )

    # Default cost entries for known provider/model combinations.
    # Costs are approximate and based on published pricing as of 2025.
    _DEFAULT_ENTRIES: ClassVar[dict[str, dict[str, CostEntry]]] = {
        "openai": {
            "gpt-4o": CostEntry(
                input_cost_per_token=0.0000025,
                output_cost_per_token=0.00001,
            ),
            "gpt-4o-mini": CostEntry(
                input_cost_per_token=0.00000015,
                output_cost_per_token=0.0000006,
            ),
            "gpt-4-turbo": CostEntry(
                input_cost_per_token=0.00001,
                output_cost_per_token=0.00003,
            ),
        },
        "anthropic": {
            "claude-sonnet-4-20250514": CostEntry(
                input_cost_per_token=0.000003,
                output_cost_per_token=0.000015,
            ),
            "claude-3-5-haiku-20241022": CostEntry(
                input_cost_per_token=0.0000008,
                output_cost_per_token=0.000004,
            ),
            "claude-opus-4-20250514": CostEntry(
                input_cost_per_token=0.000015,
                output_cost_per_token=0.000075,
            ),
        },
        "google": {
            "gemini-2.0-flash": CostEntry(
                input_cost_per_token=0.0000001,
                output_cost_per_token=0.0000004,
            ),
            "gemini-1.5-pro": CostEntry(
                input_cost_per_token=0.00000125,
                output_cost_per_token=0.000005,
            ),
        },
        "groq": {
            "llama-3.1-70b-versatile": CostEntry(
                input_cost_per_token=0.00000059,
                output_cost_per_token=0.00000079,
            ),
        },
        "together": {
            "meta-llama/Llama-3.1-70B-Instruct-Turbo": CostEntry(
                input_cost_per_token=0.00000088,
                output_cost_per_token=0.00000088,
            ),
        },
        "cohere": {
            "command-r-plus": CostEntry(
                input_cost_per_token=0.0000025,
                output_cost_per_token=0.00001,
            ),
        },
        "mock": {
            "mock-gpt": CostEntry(
                input_cost_per_token=0.0,
                output_cost_per_token=0.0,
            ),
        },
    }

    def __init__(
        self,
        custom_entries: dict[str, dict[str, CostEntry]] | None = None,
    ) -> None:
        """Initialize the cost table.

        Args:
            custom_entries: Optional custom cost entries to merge with
                defaults. Custom entries take precedence.

        Raises:
            TypeError: If custom_entries is not a dict.
        """
        if custom_entries is not None and not isinstance(custom_entries, dict):
            raise TypeError(
                f"custom_entries must be a dict, got {type(custom_entries)}"
            )
        self._entries: dict[str, dict[str, CostEntry]] = {}
        for provider, models in self._DEFAULT_ENTRIES.items():
            self._entries[provider] = dict(models)
        if custom_entries:
            for provider, models in custom_entries.items():
                if provider not in self._entries:
                    self._entries[provider] = {}
                self._entries[provider].update(models)

    def get_cost_entry(self, provider: str, model: str) -> CostEntry:
        """Look up the cost entry for a provider/model pair.

        Args:
            provider: Provider name (e.g. 'openai', 'anthropic').
            model: Model identifier (e.g. 'gpt-4o').

        Returns:
            The matching CostEntry, or DEFAULT_COST if not found.

        Raises:
            ValueError: If provider or model is empty.
        """
        if not provider:
            raise ValueError("provider must be a non-empty string")
        if not model:
            raise ValueError("model must be a non-empty string")
        provider_lower = provider.strip().lower()
        model_lower = model.strip().lower()
        provider_models = self._entries.get(provider_lower, {})
        for key, entry in provider_models.items():
            if key.lower() == model_lower:
                return entry
        return self.DEFAULT_COST

    def list_providers(self) -> list[str]:
        """Return sorted list of providers with cost data."""
        return sorted(self._entries.keys())

    def list_models(self, provider: str) -> list[str]:
        """Return sorted list of models for a provider.

        Args:
            provider: Provider name.

        Returns:
            List of model names, empty if provider unknown.

        Raises:
            ValueError: If provider is empty.
        """
        if not provider:
            raise ValueError("provider must be a non-empty string")
        return sorted(self._entries.get(provider.lower(), {}).keys())
