from typing import Optional, Dict, Any, List

# Pricing information for different providers and models
PRICING = {
    "openai": {
        "gpt-4o": {
            "input_cost": 0.01,  # $ per 1K tokens
            "output_cost": 0.03,  # $ per 1K tokens
            "context_window": 128000,
            "features": ["vision", "multilingual", "function_calling"]
        },
        "gpt-4o-mini": {
            "input_cost": 0.00075,
            "output_cost": 0.00225,
            "context_window": 128000,
            "features": ["fast", "cost-effective", "function_calling"]
        },
        "gpt-3.5-turbo": {
            "input_cost": 0.0015,
            "output_cost": 0.002,
            "context_window": 16385,
            "features": ["fast", "balanced"]
        },
        "gpt-4": {
            "input_cost": 0.03,
            "output_cost": 0.06,
            "context_window": 8192,
            "features": ["high-quality", "reasoning"]
        }
    },
    "anthropic": {
        "claude-3-5-sonnet-20241022": {
            "input_cost": 0.00_25,
            "output_cost": 0.00_75,
            "context_window": 200_000,
            "features": ["coding", "reasoning", "multilingual"]
        },
        "claude-3-5-haiku-20241022": {
            "input_cost": 0.00_00_25,
            "output_cost": 0.00_00_75,
            "context_window": 200_000,
            "features": ["fast", "cost-effective"]
        },
        "claude-3-opus-20240229": {
            "input_cost": 0.00_15,
            "output_cost": 0.00_75,
            "context_window": 200_000,
            "features": ["highest-quality", "complex-reasoning"]
        }
    },
    "cohere": {
        "command": {
            "input_cost": 0.0_12,
            "output_cost": 0.0_12,
            "context_window": 8192,
            "features": ["instruction-following", "accuracy"]
        },
        "command-light": {
            "input_cost": 0.0_04,
            "output_cost": 0.0_04,
            "context_window": 8192,
            "features": ["faster", "cheaper"]
        }
    },
    "huggingface": {
        "distilbert-base-uncased": {
            "input_cost": 0.00_01,
            "output_cost": 0.00_01,
            "context_window": 512,
            "features": ["lightweight", "fast"]
        },
        "bert-base-uncased": {
            "input_cost": 0.00_02,
            "output_cost": 0.00_02,
            "context_window": 512,
            "features": ["standard", "reliable"]
        }
    }
}


def get_pricing(
    provider: str,
    model: str
) -> Optional[Dict[str, Any]]:
    """Get pricing information for a specific provider and model"""
    return PRICING.get(provider, {}).get(model)


def calculate_cost(
    provider: str,
    model: str,
    input_tokens: int,
    output_tokens: int
) -> Dict[str, float]:
    """Calculate the cost of a request"""
    pricing = get_pricing(provider, model)
    if not pricing:
        return {"error": f"Pricing not found for {provider}/{model}"}

    input_cost = (input_tokens / 1000) * pricing["input_cost"]
    output_cost = (output_tokens / 1000) * pricing["output_cost"]
    total_cost = input_cost + output_cost

    return {
        "input_cost": input_cost,
        "output_cost": output_cost,
        "total_cost": total_cost,
        "input_tokens_per_dollar": 1000 / pricing["input_cost"] if pricing["input_cost"] > 0 else float("inf"),
        "output_tokens_per_dollar": 1000 / pricing["output_cost"] if pricing["output_cost"] > 0 else float("inf"),
        "model": model,
        "provider": provider
    }


def get_cheapest_provider(model_type: str) -> Optional[str]:
    """Get the cheapest provider for a given model type"""
    cheapest = None
    lowest_cost = float("inf")

    for provider, models in PRICING.items():
        for model, pricing in models.items():
            if model_type in model.lower():
                # Use average of input and output costs as metric
                avg_cost = (pricing["input_cost"] + pricing["output_cost"]) / 2
                if avg_cost < lowest_cost:
                    lowest_cost = avg_cost
                    cheapest = provider

    return cheapest


def compare_providers(model: str) -> List[Dict[str, Any]]:
    """Compare pricing across providers for a specific model"""
    comparisons = []

    for provider, models in PRICING.items():
        if model in models:
            pricing = models[model]
            comparisons.append({
                "provider": provider,
                "input_cost": pricing["input_cost"],
                "output_cost": pricing["output_cost"],
                "total_cost": pricing["input_cost"] + pricing["output_cost"],
                "context_window": pricing["context_window"],
                "features": pricing["features"]
            })

    return sorted(comparisons, key=lambda x: x["total_cost"])


def estimate_tokens(
    text: str,
    is_input: bool = True
) -> int:
    """Estimate token count using a simple heuristic"""
    # Rough estimate: 4 characters per token for English text
    tokens = len(text) / 4
    # Add some overhead for special tokens
    if is_input:
        tokens += 10  # Input tokens overhead
    else:
        tokens += 5   # Output tokens overhead
    return int(tokens)


def format_cost(cost: float) -> str:
    """Format cost in a human-readable way"""
    if cost >= 1:
        return f"${cost:.2f}"
    elif cost >= 0.01:
        return f"${cost:.3f}"
    else:
        return f"${cost:.4f}"


def get_cost_analysis(
    provider: str,
    model: str,
    input_text: str,
    output_text: str
) -> Dict[str, Any]:
    """Get detailed cost analysis for a request"""
    input_tokens = estimate_tokens(input_text, is_input=True)
    output_tokens = estimate_tokens(output_text, is_input=False)

    cost_info = calculate_cost(provider, model, input_tokens, output_tokens)

    return {
        "provider": provider,
        "model": model,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "cost": cost_info,
        "pricing": get_pricing(provider, model),
        "token_efficiency": {
            "input_tokens_per_dollar": cost_info.get("input_tokens_per_dollar"),
            "output_tokens_per_dollar": cost_info.get("output_tokens_per_dollar"),
        }
    }