"""
Shared pricing data and cost calculation for all integrations.

Used by langchain.py, crewai.py, autogen.py, and any future integrations.
Pricing is per 1K tokens. Updated Jan 2026.
"""

from typing import Tuple

# Format: {model_prefix: (input_cost_per_1k, output_cost_per_1k)}
MODEL_PRICING = {
    # OpenAI
    "gpt-5.2": (0.005, 0.015),
    "gpt-5": (0.004, 0.012),
    "gpt-4.1-mini": (0.0004, 0.0016),
    "gpt-4.1-nano": (0.0001, 0.0004),
    "gpt-4.1": (0.002, 0.008),
    "gpt-4o-mini": (0.00015, 0.0006),
    "gpt-4o": (0.0025, 0.01),
    "gpt-4-turbo": (0.01, 0.03),
    "gpt-4": (0.03, 0.06),
    "gpt-3.5-turbo": (0.0005, 0.0015),
    "o4-mini": (0.0011, 0.0044),
    "o3-pro": (0.02, 0.08),
    "o3-mini": (0.0011, 0.0044),
    "o3": (0.01, 0.04),
    "o1-preview": (0.015, 0.06),
    "o1-mini": (0.003, 0.012),
    "o1": (0.015, 0.06),
    # Anthropic
    "claude-opus-4": (0.015, 0.075),
    "claude-sonnet-4": (0.003, 0.015),
    "claude-4.5-opus": (0.02, 0.1),
    "claude-4.5-sonnet": (0.004, 0.02),
    "claude-4-opus": (0.018, 0.09),
    "claude-4-sonnet": (0.0035, 0.0175),
    "claude-3-7-sonnet": (0.003, 0.015),
    "claude-3-5-sonnet": (0.003, 0.015),
    "claude-3-5-haiku": (0.0008, 0.004),
    "claude-3-opus": (0.015, 0.075),
    "claude-3-sonnet": (0.003, 0.015),
    "claude-3-haiku": (0.00025, 0.00125),
    # Google Gemini 3 - Preview (Jan 2026)
    "gemini-3-pro": (0.002, 0.012),
    "gemini-3-flash": (0.0005, 0.003),
    # Google Gemini 2.5
    "gemini-2.5-pro": (0.00125, 0.005),
    "gemini-2.5-flash": (0.00015, 0.0006),
    # Google Gemini 2.0
    "gemini-2.0-flash-lite": (0.000075, 0.0003),
    "gemini-2.0-flash": (0.0001, 0.0004),
    # Google Gemini 1.5
    "gemini-1.5-pro": (0.00125, 0.005),
    "gemini-1.5-flash": (0.000075, 0.0003),
    # Google Gemini 1.0
    "gemini-1.0-pro": (0.0005, 0.0015),
    # Default fallback
    "default": (0.001, 0.002),
}

# Pre-sorted keys by length descending for correct prefix matching.
# Ensures "o3-mini" matches before "o3", "gpt-4o-mini" before "gpt-4o", etc.
_SORTED_KEYS = sorted(
    [k for k in MODEL_PRICING if k != "default"],
    key=len,
    reverse=True,
)


def get_model_pricing(model_name: str) -> Tuple[float, float]:
    """Get pricing for a model. Returns (input_cost_per_1k, output_cost_per_1k)."""
    if not model_name:
        return MODEL_PRICING["default"]

    model_lower = model_name.lower()

    for prefix in _SORTED_KEYS:
        if model_lower.startswith(prefix):
            return MODEL_PRICING[prefix]

    return MODEL_PRICING["default"]


def calculate_cost(model_name: str, prompt_tokens: int, completion_tokens: int) -> float:
    """Calculate actual cost based on token usage and model pricing."""
    input_price, output_price = get_model_pricing(model_name)

    input_cost = (prompt_tokens / 1000) * input_price
    output_cost = (completion_tokens / 1000) * output_price

    return round(input_cost + output_cost, 6)
