MODEL_PRICING: dict[str, dict[str, float]] = {
    "gemini-2.5-flash":       {"input": 0.075,  "output": 0.30},
    "gemini-2.5-flash-lite":  {"input": 0.025,  "output": 0.10},
    "gemini-2.5-pro":         {"input": 1.25,   "output": 10.0},
    "gemini-2.0-flash":       {"input": 0.10,   "output": 0.40},
    "gemini-2.0-flash-lite":  {"input": 0.025,  "output": 0.10},
    "gemini-3-flash-preview":  {"input": 0.075,  "output": 0.30},
    "gemini-3-pro-preview":    {"input": 1.25,   "output": 10.0},
    "gemini-flash-latest":    {"input": 0.075,  "output": 0.30},
    "gemini-flash-lite-latest": {"input": 0.025, "output": 0.10},
    "gemini-pro-latest":      {"input": 1.25,   "output": 10.0},
}

_DEFAULT_PRICING: dict[str, float] = {"input": 0.10, "output": 0.40}


def estimate_cost(model_name: str, input_tokens: int, output_tokens: int) -> float:
    """Return the estimated cost in USD for a given token usage.

    Args:
        model_name (str): The Gemini model identifier.
        input_tokens (int): Total prompt/input tokens consumed.
        output_tokens (int): Total generated/output tokens consumed.

    Returns:
        float: Estimated cost in USD, rounded to six decimal places.
    """
    pricing = MODEL_PRICING.get(model_name, _DEFAULT_PRICING)
    cost = (input_tokens * pricing["input"] + output_tokens * pricing["output"]) / 1_000_000
    return round(cost, 6)