"""
Pricing - Centralized cost estimation for all LLM providers.

This module provides a single source of truth for model pricing,
eliminating duplicate pricing tables across adapters.

Usage:
    from agent_orchestrator.adapters.pricing import estimate_cost, get_model_pricing

    cost = estimate_cost("claude-sonnet-4-20250514", input_tokens=1000, output_tokens=500)
    input_price, output_price = get_model_pricing("gpt-4o")
"""

from dataclasses import dataclass
from typing import Optional, Tuple


@dataclass
class ModelPricing:
    """Pricing information for a model."""

    input_price_per_1k: float  # Price per 1K input tokens
    output_price_per_1k: float  # Price per 1K output tokens
    provider: str  # Provider name (anthropic, openai, google)
    notes: str = ""  # Optional notes about pricing


# Centralized pricing table (prices as of early 2025)
# Update this single location when prices change
MODEL_PRICING: dict[str, ModelPricing] = {
    # Anthropic Claude models
    "claude-opus-4-20250514": ModelPricing(0.015, 0.075, "anthropic"),
    "claude-sonnet-4-20250514": ModelPricing(0.003, 0.015, "anthropic"),
    "claude-3-5-sonnet-20241022": ModelPricing(0.003, 0.015, "anthropic"),
    "claude-3-opus-20240229": ModelPricing(0.015, 0.075, "anthropic"),
    "claude-3-sonnet-20240229": ModelPricing(0.003, 0.015, "anthropic"),
    "claude-3-haiku-20240307": ModelPricing(0.00025, 0.00125, "anthropic"),
    "claude-3-5-haiku-20241022": ModelPricing(0.001, 0.005, "anthropic"),
    # CLI agents use default pricing
    "claude-code": ModelPricing(0.003, 0.015, "anthropic", "Claude Code CLI default"),
    # OpenAI models
    "gpt-4o": ModelPricing(0.005, 0.015, "openai"),
    "gpt-4o-mini": ModelPricing(0.00015, 0.0006, "openai"),
    "gpt-4-turbo": ModelPricing(0.01, 0.03, "openai"),
    "gpt-4": ModelPricing(0.03, 0.06, "openai"),
    "gpt-3.5-turbo": ModelPricing(0.0005, 0.0015, "openai"),
    "o1-preview": ModelPricing(0.015, 0.06, "openai"),
    "o1-mini": ModelPricing(0.003, 0.012, "openai"),
    "o1": ModelPricing(0.015, 0.06, "openai"),
    "o3-mini": ModelPricing(0.0011, 0.0044, "openai"),
    # Google Gemini models
    "gemini-pro": ModelPricing(0.0005, 0.0015, "google"),
    "gemini-1.5-pro": ModelPricing(0.00125, 0.005, "google"),
    "gemini-1.5-flash": ModelPricing(0.000075, 0.0003, "google"),
    "gemini-cli": ModelPricing(0.0005, 0.0015, "google", "Gemini CLI default"),
}

# Default pricing by provider (fallback when model not found)
DEFAULT_PRICING: dict[str, ModelPricing] = {
    "anthropic": ModelPricing(0.003, 0.015, "anthropic", "Default Claude pricing"),
    "openai": ModelPricing(0.005, 0.015, "openai", "Default GPT-4o pricing"),
    "google": ModelPricing(0.0005, 0.0015, "google", "Default Gemini pricing"),
}


def get_model_pricing(model: str, provider: Optional[str] = None) -> Tuple[float, float]:
    """
    Get pricing for a specific model.

    Args:
        model: Model name/ID
        provider: Optional provider hint for fallback

    Returns:
        Tuple of (input_price_per_1k, output_price_per_1k)
    """
    # Direct lookup
    if model in MODEL_PRICING:
        pricing = MODEL_PRICING[model]
        return (pricing.input_price_per_1k, pricing.output_price_per_1k)

    # Try to infer provider from model name
    if provider is None:
        model_lower = model.lower()
        if "claude" in model_lower:
            provider = "anthropic"
        elif "gpt" in model_lower or model_lower.startswith("o1") or model_lower.startswith("o3"):
            provider = "openai"
        elif "gemini" in model_lower:
            provider = "google"

    # Use provider default
    if provider and provider in DEFAULT_PRICING:
        pricing = DEFAULT_PRICING[provider]
        return (pricing.input_price_per_1k, pricing.output_price_per_1k)

    # Ultimate fallback (Sonnet pricing)
    return (0.003, 0.015)


def estimate_cost(
    model: str,
    input_tokens: int,
    output_tokens: int,
    provider: Optional[str] = None,
) -> float:
    """
    Estimate cost for a request.

    Args:
        model: Model name/ID
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        provider: Optional provider hint

    Returns:
        Estimated cost in USD
    """
    input_price, output_price = get_model_pricing(model, provider)
    return (input_tokens / 1000 * input_price) + (output_tokens / 1000 * output_price)


def get_provider_models(provider: str) -> list[str]:
    """Get all models for a specific provider."""
    return [
        model
        for model, pricing in MODEL_PRICING.items()
        if pricing.provider == provider
    ]


def get_all_models() -> list[str]:
    """Get all known models."""
    return list(MODEL_PRICING.keys())


def add_model_pricing(
    model: str,
    input_price_per_1k: float,
    output_price_per_1k: float,
    provider: str,
    notes: str = "",
) -> None:
    """
    Add or update pricing for a model.

    Useful for runtime configuration of custom models.

    Args:
        model: Model name/ID
        input_price_per_1k: Price per 1K input tokens
        output_price_per_1k: Price per 1K output tokens
        provider: Provider name
        notes: Optional notes
    """
    MODEL_PRICING[model] = ModelPricing(
        input_price_per_1k,
        output_price_per_1k,
        provider,
        notes,
    )
