"""Per-model token pricing table and cost calculation."""

from __future__ import annotations

# model_id_prefix → (input_price_per_1k_tokens, output_price_per_1k_tokens)
PRICING: dict[str, tuple[float, float]] = {
    "us.anthropic.claude-sonnet-4": (0.003, 0.015),
    "us.anthropic.claude-opus-4": (0.015, 0.075),
    "us.anthropic.claude-haiku-4": (0.0008, 0.004),
    "amazon.nova-pro": (0.0008, 0.0032),
    "amazon.nova-lite": (0.00006, 0.00024),
    "amazon.nova-micro": (0.000035, 0.00014),
    "us.deepseek.r1": (0.00135, 0.00539),
}


def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
    """Compute approximate cost in USD.

    Looks up the model in the PRICING table by prefix match and applies:
        (input_tokens / 1000) * input_rate + (output_tokens / 1000) * output_rate

    Returns 0.0 for unknown models.
    """
    for prefix, (input_rate, output_rate) in PRICING.items():
        if model.startswith(prefix):
            return (input_tokens / 1000) * input_rate + (output_tokens / 1000) * output_rate
    return 0.0


def format_cost(cost: float) -> str:
    """Format a cost value as '$X.XXXX' string."""
    return f"${cost:.4f}"
