"""
Utility functions for Visibe SDK.
"""

# Model pricing per 1K tokens (in USD)
# Sorted by key length descending for longest-prefix matching.
MODEL_PRICING = {
    # OpenAI
    "gpt-4": {"input": 0.03, "output": 0.06},
    "gpt-4-turbo": {"input": 0.01, "output": 0.03},
    "gpt-4-turbo-preview": {"input": 0.01, "output": 0.03},
    "gpt-4o": {"input": 0.0025, "output": 0.01},
    "gpt-4o-mini": {"input": 0.00015, "output": 0.0006},
    "gpt-4.1": {"input": 0.002, "output": 0.008},
    "gpt-4.1-mini": {"input": 0.0004, "output": 0.0016},
    "gpt-4.1-nano": {"input": 0.0001, "output": 0.0004},
    "gpt-3.5-turbo": {"input": 0.0015, "output": 0.002},
    "o1": {"input": 0.015, "output": 0.06},
    "o1-mini": {"input": 0.003, "output": 0.012},
    "o3-mini": {"input": 0.0011, "output": 0.0044},
    # Anthropic (via OpenAI-compatible proxies)
    "claude-3-opus": {"input": 0.015, "output": 0.075},
    "claude-3-sonnet": {"input": 0.003, "output": 0.015},
    "claude-3-haiku": {"input": 0.00025, "output": 0.00125},
    "claude-3.5-sonnet": {"input": 0.003, "output": 0.015},
    "claude-3.5-haiku": {"input": 0.0008, "output": 0.004},
    # Google (via OpenAI-compatible proxies)
    "gemini-1.5-pro": {"input": 0.00125, "output": 0.005},
    "gemini-1.5-flash": {"input": 0.000075, "output": 0.0003},
    "gemini-2.0-flash": {"input": 0.0001, "output": 0.0004},
    # DeepSeek (via OpenAI-compatible proxies)
    "deepseek-chat": {"input": 0.00027, "output": 0.0011},
    "deepseek-reasoner": {"input": 0.00055, "output": 0.00219},
    # AWS Bedrock — Anthropic
    "anthropic.claude-3-opus": {"input": 0.015, "output": 0.075},
    "anthropic.claude-3-sonnet": {"input": 0.003, "output": 0.015},
    "anthropic.claude-3-haiku": {"input": 0.00025, "output": 0.00125},
    "anthropic.claude-3.5-sonnet": {"input": 0.003, "output": 0.015},
    "anthropic.claude-3.5-haiku": {"input": 0.0008, "output": 0.004},
    "anthropic.claude-4-sonnet": {"input": 0.003, "output": 0.015},
    "anthropic.claude-4-opus": {"input": 0.015, "output": 0.075},
    # AWS Bedrock — Meta
    "meta.llama3-70b-instruct": {"input": 0.00265, "output": 0.0035},
    "meta.llama3-8b-instruct": {"input": 0.0003, "output": 0.0006},
    "meta.llama3-1-405b-instruct": {"input": 0.00532, "output": 0.016},
    "meta.llama3-1-70b-instruct": {"input": 0.00265, "output": 0.0035},
    "meta.llama3-1-8b-instruct": {"input": 0.0003, "output": 0.0006},
    # AWS Bedrock — Mistral
    "mistral.mistral-large": {"input": 0.004, "output": 0.012},
    "mistral.mistral-small": {"input": 0.001, "output": 0.003},
    "mistral.mixtral-8x7b-instruct": {"input": 0.00045, "output": 0.0007},
    # AWS Bedrock — Amazon
    "amazon.titan-text-express": {"input": 0.0002, "output": 0.0006},
    "amazon.titan-text-lite": {"input": 0.00015, "output": 0.0002},
    "amazon.nova-pro": {"input": 0.0008, "output": 0.0032},
    "amazon.nova-lite": {"input": 0.00006, "output": 0.00024},
    "amazon.nova-micro": {"input": 0.000035, "output": 0.00014},
    # AWS Bedrock — Cohere
    "cohere.command-r-plus": {"input": 0.003, "output": 0.015},
    "cohere.command-r": {"input": 0.0005, "output": 0.0015},
    # AWS Bedrock — AI21
    "ai21.jamba-1.5-large": {"input": 0.002, "output": 0.008},
    "ai21.jamba-1.5-mini": {"input": 0.0002, "output": 0.0004},
}


def calculate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    """
    Calculate cost based on model and token usage.

    Args:
        model: Model name (e.g., "gpt-4", "gpt-3.5-turbo")
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens

    Returns:
        Cost in USD
    """
    # Normalize model name
    model_lower = model.lower()
    model_base = None

    # Sort by length descending so "gpt-4o-mini" matches before "gpt-4"
    for known_model in sorted(MODEL_PRICING.keys(), key=len, reverse=True):
        if known_model in model_lower:
            model_base = known_model
            break

    if not model_base:
        return 0.0

    pricing = MODEL_PRICING[model_base]
    input_cost = (input_tokens / 1000) * pricing['input']
    output_cost = (output_tokens / 1000) * pricing['output']

    return round(input_cost + output_cost, 6)
