"""
Unified Pricing Database for AgentBill
Prevents cost bleeding by centralizing all model pricing
"""

from typing import Dict, Optional

# Pricing per 1M tokens (input, output)
MODEL_PRICING: Dict[str, Dict[str, float]] = {
    # OpenAI Chat Models
    "gpt-5": {"input": 2.50, "output": 10.00},
    "gpt-5-mini": {"input": 0.15, "output": 0.60},
    "gpt-5-nano": {"input": 0.10, "output": 0.40},
    "gpt-4o": {"input": 5.00, "output": 15.00},
    "gpt-4o-mini": {"input": 0.15, "output": 0.60},
    "gpt-4": {"input": 30.00, "output": 60.00},
    "gpt-3.5-turbo": {"input": 0.50, "output": 1.50},
    
    # OpenAI Embeddings (per 1M tokens, no output)
    "text-embedding-3-small": {"input": 0.02, "output": 0},
    "text-embedding-3-large": {"input": 0.13, "output": 0},
    "text-embedding-ada-002": {"input": 0.10, "output": 0},
    
    # Anthropic
    "claude-sonnet-4-5": {"input": 3.00, "output": 15.00},
    "claude-opus-4-1-20250805": {"input": 15.00, "output": 75.00},
    "claude-3-7-sonnet-20250219": {"input": 3.00, "output": 15.00},
    "claude-3-5-sonnet-20241022": {"input": 3.00, "output": 15.00},
    "claude-3-5-haiku-20241022": {"input": 0.80, "output": 4.00},
    
    # Google Gemini
    "gemini-2.5-pro": {"input": 1.25, "output": 5.00},
    "gemini-2.5-flash": {"input": 0.075, "output": 0.30},
    "gemini-2.5-flash-lite": {"input": 0.019, "output": 0.075},
    "gemini-1.5-pro": {"input": 1.25, "output": 5.00},
    "gemini-pro": {"input": 0.50, "output": 1.50},  # Legacy
    "text-embedding-004": {"input": 0.025, "output": 0},
    
    # Mistral
    "mistral-large-latest": {"input": 4.00, "output": 12.00},
    "mistral-medium-latest": {"input": 2.70, "output": 8.10},
    "mistral-small-latest": {"input": 1.00, "output": 3.00},
    "mistral-embed": {"input": 0.10, "output": 0},
    
    # Perplexity (NEW)
    "llama-3.1-sonar-huge-128k-online": {"input": 5.00, "output": 5.00},
    "llama-3.1-sonar-large-128k-online": {"input": 1.00, "output": 1.00},
    "llama-3.1-sonar-small-128k-online": {"input": 0.20, "output": 0.20},
    "llama-3.1-8b-instruct": {"input": 0.20, "output": 0.20},
    "llama-3.1-70b-instruct": {"input": 1.00, "output": 1.00},
    
    # Ollama (Local - zero cost, but track metrics)
    "ollama": {"input": 0.00, "output": 0.00},  # All Ollama models
    
    # AWS Bedrock
    "anthropic.claude-3-5-sonnet": {"input": 3.00, "output": 15.00},
    "anthropic.claude-3-haiku": {"input": 0.25, "output": 1.25},
    "amazon.titan-text-express": {"input": 0.20, "output": 0.60},
}

# Special pricing for non-token-based operations
OPERATION_PRICING = {
    # OpenAI Images (per image)
    "dall-e-3-standard-1024": 0.040,
    "dall-e-3-hd-1024": 0.080,
    "dall-e-2-1024": 0.020,
    
    # OpenAI Audio (per minute for Whisper)
    "whisper-1": 0.006,  # per minute
    
    # OpenAI TTS (per 1M characters)
    "tts-1": 15.00,
    "tts-1-hd": 30.00,
}


def calculate_cost(
    model: str,
    input_tokens: int,
    output_tokens: int,
    provider: str = "openai"
) -> float:
    """
    Calculate cost for a given model and token usage.
    
    Args:
        model: Model name (e.g., 'gpt-4o', 'claude-sonnet-4-5')
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        provider: Provider name for context
        
    Returns:
        Cost in USD
    """
    # Normalize model name
    model_key = model.lower()
    
    # Handle Ollama special case
    if provider == "ollama" or model_key.startswith("ollama/"):
        return 0.0  # Local execution, no cost
    
    # Get pricing or use conservative fallback
    pricing = MODEL_PRICING.get(model_key, {"input": 1.00, "output": 2.00})
    
    # Calculate cost (pricing is per 1M tokens)
    cost = (input_tokens / 1_000_000 * pricing["input"]) + \
           (output_tokens / 1_000_000 * pricing["output"])
    
    return cost


def calculate_image_cost(model: str, size: str, quality: str = "standard") -> float:
    """Calculate cost for image generation"""
    key = f"{model}-{quality}-{size.split('x')[0]}"
    return OPERATION_PRICING.get(key, 0.04)  # Default $0.04


def calculate_audio_cost(model: str, duration_seconds: float = 0, chars: int = 0) -> float:
    """
    Calculate cost for audio operations.
    
    For Whisper: duration_seconds
    For TTS: chars (number of characters)
    """
    if model == "whisper-1":
        minutes = duration_seconds / 60.0
        return minutes * OPERATION_PRICING.get("whisper-1", 0.006)
    elif model.startswith("tts-"):
        return (chars / 1_000_000) * OPERATION_PRICING.get(model, 15.00)
    return 0.0


def is_embedding_model(model: str) -> bool:
    """Check if model is an embedding model"""
    return "embedding" in model.lower() or "embed" in model.lower()


def get_provider_from_model(model: str) -> str:
    """Infer provider from model name"""
    model_lower = model.lower()
    
    if model_lower.startswith("gpt-") or "dall-e" in model_lower or "whisper" in model_lower or "tts" in model_lower:
        return "openai"
    elif model_lower.startswith("claude-"):
        return "anthropic"
    elif model_lower.startswith("gemini-") or "embedding-004" in model_lower:
        return "google"
    elif model_lower.startswith("mistral-"):
        return "mistral"
    elif "sonar" in model_lower or model_lower.startswith("llama-3.1-"):
        return "perplexity"
    elif model_lower.startswith("ollama/") or "ollama" in model_lower:
        return "ollama"
    elif model_lower.startswith("anthropic.") or model_lower.startswith("amazon."):
        return "bedrock"
    else:
        return "unknown"
