import json
import os
import sys

import pytest
from fastapi.testclient import TestClient

import litellm
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
    StandardBuiltInToolCostTracking,
)
from litellm.llms.gemini.image_generation.cost_calculator import (
    cost_calculator as gemini_image_generation_cost_calculator,
)
from litellm.llms.vertex_ai.image_generation.cost_calculator import (
    cost_calculator as vertex_image_generation_cost_calculator,
)
from litellm.types.llms.openai import FileSearchTool, WebSearchOptions
from litellm.types.utils import (
    CompletionTokensDetailsWrapper,
    ImageObject,
    ImageResponse,
    ImageUsage,
    ImageUsageInputTokensDetails,
    ModelInfo,
    ModelResponse,
    PromptTokensDetailsWrapper,
    StandardBuiltInToolsParams,
)

sys.path.insert(
    0, os.path.abspath("../../..")
)  # Adds the parent directory to the system path

from litellm.litellm_core_utils.llm_cost_calc.utils import (
    PromptTokensDetailsResult,
    _calculate_input_cost,
    calculate_cache_writing_cost,
    generic_cost_per_token,
)
from litellm.types.utils import CacheCreationTokenDetails, Usage


def test_reasoning_tokens_no_price_set():
    # Use o1 - o1-mini was deprecated/renamed; o1 has same reasoning-token semantics
    # (no separate output_cost_per_reasoning_token, so all completion tokens use output_cost_per_token)
    model = "o1"
    custom_llm_provider = "openai"
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")
    model_cost_map = litellm.model_cost[model]
    usage = Usage(
        completion_tokens=1578,
        prompt_tokens=17,
        total_tokens=1595,
        completion_tokens_details=CompletionTokensDetailsWrapper(
            accepted_prediction_tokens=None,
            audio_tokens=None,
            reasoning_tokens=952,
            rejected_prediction_tokens=None,
            text_tokens=626,
        ),
        prompt_tokens_details=PromptTokensDetailsWrapper(
            audio_tokens=None, cached_tokens=None, text_tokens=17, image_tokens=None
        ),
    )
    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider="openai",
    )
    assert round(prompt_cost, 10) == round(
        model_cost_map["input_cost_per_token"] * usage.prompt_tokens,
        10,
    )
    print(f"completion_cost: {completion_cost}")
    expected_completion_cost = (
        model_cost_map["output_cost_per_token"] * usage.completion_tokens
    )
    print(f"expected_completion_cost: {expected_completion_cost}")
    assert round(completion_cost, 10) == round(
        expected_completion_cost,
        10,
    )


def test_reasoning_tokens_gemini():
    model = "gemini-2.5-flash"
    custom_llm_provider = "gemini"
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    usage = Usage(
        completion_tokens=1578,
        prompt_tokens=17,
        total_tokens=1595,
        completion_tokens_details=CompletionTokensDetailsWrapper(
            accepted_prediction_tokens=None,
            audio_tokens=None,
            reasoning_tokens=952,
            rejected_prediction_tokens=None,
            text_tokens=626,
        ),
        prompt_tokens_details=PromptTokensDetailsWrapper(
            audio_tokens=None, cached_tokens=None, text_tokens=17, image_tokens=None
        ),
    )
    model_cost_map = litellm.model_cost[model]
    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )

    assert round(prompt_cost, 10) == round(
        model_cost_map["input_cost_per_token"] * usage.prompt_tokens,
        10,
    )
    assert round(completion_cost, 10) == round(
        (
            model_cost_map["output_cost_per_token"]
            * usage.completion_tokens_details.text_tokens
        )
        + (
            model_cost_map["output_cost_per_reasoning_token"]
            * usage.completion_tokens_details.reasoning_tokens
        ),
        10,
    )


def test_reasoning_tokens_gemini_3_1_flash_lite():
    """Test cost calculation for gemini-3.1-flash-lite-preview with reasoning tokens"""
    model = "gemini-3.1-flash-lite-preview"
    custom_llm_provider = "gemini"
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    usage = Usage(
        completion_tokens=1000,
        prompt_tokens=500,
        total_tokens=1500,
        completion_tokens_details=CompletionTokensDetailsWrapper(
            accepted_prediction_tokens=None,
            audio_tokens=None,
            reasoning_tokens=400,
            rejected_prediction_tokens=None,
            text_tokens=600,
        ),
        prompt_tokens_details=PromptTokensDetailsWrapper(
            audio_tokens=None, cached_tokens=None, text_tokens=500, image_tokens=None
        ),
    )
    model_cost_map = litellm.model_cost[model]
    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )

    assert round(prompt_cost, 10) == round(
        model_cost_map["input_cost_per_token"] * usage.prompt_tokens,
        10,
    )
    assert round(completion_cost, 10) == round(
        (
            model_cost_map["output_cost_per_token"]
            * usage.completion_tokens_details.text_tokens
        )
        + (
            model_cost_map["output_cost_per_reasoning_token"]
            * usage.completion_tokens_details.reasoning_tokens
        ),
        10,
    )


def test_image_tokens_with_custom_pricing():
    """Test that image_tokens in completion are properly costed with output_cost_per_image_token."""
    from unittest.mock import patch

    # Mock model info with image token pricing
    mock_model_info = {
        "input_cost_per_token": 1e-6,
        "output_cost_per_token": 2e-6,
        "output_cost_per_image_token": 5e-6,  # Custom pricing for image tokens in output
    }

    usage = Usage(
        completion_tokens=1720,  # text_tokens (600) + image_tokens (1120)
        prompt_tokens=14,
        total_tokens=1734,
        completion_tokens_details=CompletionTokensDetailsWrapper(
            accepted_prediction_tokens=None,
            audio_tokens=None,
            reasoning_tokens=0,
            rejected_prediction_tokens=None,
            text_tokens=600,
            image_tokens=1120,
        ),
        prompt_tokens_details=PromptTokensDetailsWrapper(
            audio_tokens=None, cached_tokens=None, text_tokens=14, image_tokens=None
        ),
    )

    with patch(
        "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
        return_value=mock_model_info,
    ):
        prompt_cost, completion_cost = generic_cost_per_token(
            model="test-model", usage=usage, custom_llm_provider="gemini"
        )

    # Expected costs:
    # Prompt: 14 * 1e-6
    # Completion: (600 * 2e-6) + (1120 * 5e-6)
    expected_prompt_cost = 14 * 1e-6
    expected_completion_cost = (600 * 2e-6) + (1120 * 5e-6)

    assert round(prompt_cost, 12) == round(expected_prompt_cost, 12)
    assert round(completion_cost, 12) == round(expected_completion_cost, 12)


def test_image_tokens_fallback_to_base_cost():
    """Test that image_tokens fall back to base cost when output_cost_per_image_token is not set."""
    from unittest.mock import patch

    # Mock model info without image token pricing
    mock_model_info = {
        "input_cost_per_token": 1e-6,
        "output_cost_per_token": 2e-6,
        # No output_cost_per_image_token defined
    }

    usage = Usage(
        completion_tokens=1720,
        prompt_tokens=14,
        total_tokens=1734,
        completion_tokens_details=CompletionTokensDetailsWrapper(
            accepted_prediction_tokens=None,
            audio_tokens=None,
            reasoning_tokens=0,
            rejected_prediction_tokens=None,
            text_tokens=600,
            image_tokens=1120,
        ),
        prompt_tokens_details=PromptTokensDetailsWrapper(
            audio_tokens=None, cached_tokens=None, text_tokens=14, image_tokens=None
        ),
    )

    with patch(
        "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
        return_value=mock_model_info,
    ):
        prompt_cost, completion_cost = generic_cost_per_token(
            model="test-model", usage=usage, custom_llm_provider="gemini"
        )

    # Expected costs:
    # Prompt: 14 * 1e-6
    # Completion: (600 * 2e-6) + (1120 * 2e-6)  # image_tokens use base cost
    expected_prompt_cost = 14 * 1e-6
    expected_completion_cost = (600 * 2e-6) + (1120 * 2e-6)

    assert round(prompt_cost, 12) == round(expected_prompt_cost, 12)
    assert round(completion_cost, 12) == round(expected_completion_cost, 12)


def test_generic_cost_per_token_above_200k_tokens():
    # gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing
    model = "gemini-2.5-pro"
    custom_llm_provider = "vertex_ai"
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    model_cost_map = litellm.model_cost[model]
    prompt_tokens = 220 * 1e6
    completion_tokens = 150
    usage = Usage(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=prompt_tokens + completion_tokens,
    )
    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )
    assert round(prompt_cost, 10) == round(
        model_cost_map["input_cost_per_token_above_200k_tokens"] * usage.prompt_tokens,
        10,
    )
    assert round(completion_cost, 10) == round(
        model_cost_map["output_cost_per_token_above_200k_tokens"]
        * usage.completion_tokens,
        10,
    )


def test_generic_cost_per_token_gpt54_above_272k_tokens():
    """GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output."""
    model = "gpt-5.4"
    custom_llm_provider = "openai"
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    model_cost_map = litellm.model_cost[model]
    prompt_tokens = 273000  # Above 272K threshold
    completion_tokens = 1000
    usage = Usage(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=prompt_tokens + completion_tokens,
    )
    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )
    expected_prompt = (
        model_cost_map["input_cost_per_token_above_272k_tokens"] * prompt_tokens
    )
    expected_completion = (
        model_cost_map["output_cost_per_token_above_272k_tokens"] * completion_tokens
    )
    assert round(prompt_cost, 10) == round(expected_prompt, 10)
    assert round(completion_cost, 10) == round(expected_completion, 10)


def test_generic_cost_per_token_minimax_m3_above_512k_tokens():
    """MiniMax-M3: prompts >512K input tokens priced at 2x input, output, and cache read."""
    model = "minimax/MiniMax-M3"
    custom_llm_provider = "minimax"
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    model_cost_map = litellm.model_cost[model]
    prompt_tokens = 600000
    cached_tokens = 100000
    completion_tokens = 1000
    usage = Usage(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=prompt_tokens + completion_tokens,
        prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached_tokens),
    )
    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )
    expected_prompt = (
        model_cost_map["input_cost_per_token_above_512k_tokens"]
        * (prompt_tokens - cached_tokens)
        + model_cost_map["cache_read_input_token_cost_above_512k_tokens"]
        * cached_tokens
    )
    expected_completion = (
        model_cost_map["output_cost_per_token_above_512k_tokens"] * completion_tokens
    )
    assert round(prompt_cost, 10) == round(expected_prompt, 10)
    assert round(completion_cost, 10) == round(expected_completion, 10)


def test_generic_cost_per_token_gpt55():
    """gpt-5.5: base pricing — $5/1M input, $30/1M output, $0.50/1M cached input."""
    model = "gpt-5.5"
    custom_llm_provider = "openai"
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    model_cost_map = litellm.model_cost[model]

    # Sanity-check the map values match OpenAI's published pricing.
    assert model_cost_map["input_cost_per_token"] == 5e-6
    assert model_cost_map["output_cost_per_token"] == 3e-5
    assert model_cost_map["cache_read_input_token_cost"] == 5e-7
    assert model_cost_map["litellm_provider"] == "openai"
    assert model_cost_map["mode"] == "chat"
    # gpt-5.5 inherits GPT-5.4's long-context window + tiered pricing.
    assert model_cost_map["max_input_tokens"] == 1050000
    assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 1e-5
    assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 4.5e-5

    prompt_tokens = 1000
    completion_tokens = 500
    usage = Usage(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=prompt_tokens + completion_tokens,
    )
    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )
    assert round(prompt_cost, 10) == round(
        model_cost_map["input_cost_per_token"] * prompt_tokens, 10
    )
    assert round(completion_cost, 10) == round(
        model_cost_map["output_cost_per_token"] * completion_tokens, 10
    )


def test_generic_cost_per_token_gpt55_pro():
    """gpt-5.5-pro: responses-only model — $30/1M input, $180/1M output, $3/1M cached input."""
    model = "gpt-5.5-pro"
    custom_llm_provider = "openai"
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    model_cost_map = litellm.model_cost[model]

    # Sanity-check the map values match OpenAI's published pricing.
    assert model_cost_map["input_cost_per_token"] == 3e-5
    assert model_cost_map["output_cost_per_token"] == 1.8e-4
    assert model_cost_map["cache_read_input_token_cost"] == 3e-6
    assert model_cost_map["litellm_provider"] == "openai"
    # gpt-5.5-pro is a responses-only model (no /v1/chat/completions endpoint).
    assert model_cost_map["mode"] == "responses"
    assert "/v1/chat/completions" not in model_cost_map["supported_endpoints"]
    assert "/v1/responses" in model_cost_map["supported_endpoints"]
    # Inherits GPT-5.4-pro's long-context window + tiered pricing.
    assert model_cost_map["max_input_tokens"] == 1050000
    assert model_cost_map["input_cost_per_token_above_272k_tokens"] == 6e-5
    assert model_cost_map["output_cost_per_token_above_272k_tokens"] == 2.7e-4

    prompt_tokens = 1000
    completion_tokens = 500
    usage = Usage(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=prompt_tokens + completion_tokens,
    )
    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )
    assert round(prompt_cost, 10) == round(
        model_cost_map["input_cost_per_token"] * prompt_tokens, 10
    )
    assert round(completion_cost, 10) == round(
        model_cost_map["output_cost_per_token"] * completion_tokens, 10
    )


@pytest.mark.parametrize(
    "model,expected_none,expected_xhigh,expected_minimal",
    [
        # Verified against OpenAI's live API on 2026-04-24:
        #   gpt-5.5   -> supports: none, low, medium, high, xhigh
        #   gpt-5.5-pro -> supports: medium, high, xhigh
        # Neither supports "minimal"; gpt-5.5-pro additionally does not support "none".
        # The JSON must reflect this so LiteLLM rejects unsupported values locally
        # (or drops them with drop_params=True) instead of round-tripping to OpenAI
        # for a 400.
        ("gpt-5.5", True, True, False),
        ("gpt-5.5-2026-04-23", True, True, False),
        ("gpt-5.5-pro", False, True, False),
        ("gpt-5.5-pro-2026-04-23", False, True, False),
    ],
)
def test_gpt55_reasoning_effort_flags_match_live_openai_api(
    model, expected_none, expected_xhigh, expected_minimal
):
    """Pin reasoning_effort capability flags to OpenAI's actual API contract.

    Observed via `POST /v1/chat/completions` with reasoning_effort=minimal:
    ``Unsupported value: 'reasoning_effort' does not support 'minimal' with
    this model``. gpt-5.5-pro additionally rejects 'none' and 'low'.
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    m = litellm.model_cost[model]
    assert (
        m.get("supports_none_reasoning_effort") is expected_none
    ), f"{model}: supports_none_reasoning_effort expected {expected_none}"
    assert (
        m.get("supports_xhigh_reasoning_effort") is expected_xhigh
    ), f"{model}: supports_xhigh_reasoning_effort expected {expected_xhigh}"
    assert (
        m.get("supports_minimal_reasoning_effort") is expected_minimal
    ), f"{model}: supports_minimal_reasoning_effort expected {expected_minimal}"


@pytest.mark.parametrize(
    "base_model,dated_model",
    [
        ("gpt-5.5", "gpt-5.5-2026-04-23"),
        ("gpt-5.5-pro", "gpt-5.5-pro-2026-04-23"),
    ],
)
def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(
    base_model, dated_model
):
    """Dated snapshots must carry the same reasoning_effort capability flags as
    their non-dated counterparts.

    Regression guard: ``supports_{none,minimal,xhigh}_reasoning_effort`` gate
    downstream routing in ``OpenAIGPT5Config`` — a missing flag is treated as
    ``False`` for opt-in levels (e.g. ``xhigh``), which silently diverges
    behavior between ``gpt-5.5`` and ``gpt-5.5-2026-04-23``. Pinning to a
    dated variant must never lose capabilities relative to the base alias.
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    base = litellm.model_cost[base_model]
    dated = litellm.model_cost[dated_model]

    for flag in (
        "supports_none_reasoning_effort",
        "supports_minimal_reasoning_effort",
        "supports_xhigh_reasoning_effort",
    ):
        assert dated.get(flag) == base.get(flag), (
            f"{dated_model} has {flag}={dated.get(flag)!r}, "
            f"but {base_model} has {flag}={base.get(flag)!r}. "
            f"Dated snapshots must inherit the base model's reasoning_effort "
            f"capability profile."
        )


@pytest.mark.parametrize(
    "model,expected_mode,expected_input,expected_output,expected_cache_read",
    [
        ("azure/gpt-5.5", "chat", 5e-6, 3e-5, 5e-7),
        ("azure/gpt-5.5-2026-04-23", "chat", 5e-6, 3e-5, 5e-7),
        ("azure/gpt-5.5-pro", "responses", 3e-5, 1.8e-4, 3e-6),
        ("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6),
    ],
)
def test_azure_gpt55_entries_present_with_correct_pricing(
    model, expected_mode, expected_input, expected_output, expected_cache_read
):
    """Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure.

    Pricing parity with openai/gpt-5.5* (verified against OpenAI's pricing page
    on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro.
    Cache discount is 10% of input.
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    m = litellm.model_cost[model]
    assert m["litellm_provider"] == "azure"
    assert m["mode"] == expected_mode
    assert m["input_cost_per_token"] == expected_input
    assert m["output_cost_per_token"] == expected_output
    assert m["cache_read_input_token_cost"] == expected_cache_read
    # Long-context window inherited from gpt-5.4 / openai gpt-5.5.
    assert m["max_input_tokens"] == 1050000
    assert m["max_output_tokens"] == 128000


@pytest.mark.parametrize(
    "model,expected_none,expected_minimal,expected_xhigh",
    [
        # Mirror live OpenAI API contract (verified via openai/gpt-5.5* on
        # 2026-04-24): chat accepts {none, low, medium, high, xhigh} but NOT
        # minimal; pro accepts {medium, high, xhigh} only.
        # NOTE: openai/gpt-5.5* entries currently set supports_minimal=true on
        # main (pre #26456). Once that PR lands, OpenAI + Azure flags align.
        ("azure/gpt-5.5", True, False, True),
        ("azure/gpt-5.5-pro", False, False, True),
    ],
)
def test_azure_gpt55_reasoning_effort_flags_match_live_openai_api(
    model, expected_none, expected_minimal, expected_xhigh
):
    """Azure entries pin reasoning_effort flags to OpenAI's actual API contract."""
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    m = litellm.model_cost[model]
    assert m.get("supports_none_reasoning_effort") is expected_none
    assert m.get("supports_minimal_reasoning_effort") is expected_minimal
    assert m.get("supports_xhigh_reasoning_effort") is expected_xhigh


def test_generic_cost_per_token_anthropic_prompt_caching():
    model = "claude-sonnet-4@20250514"
    usage = Usage(
        completion_tokens=90,
        prompt_tokens=28436,
        total_tokens=28526,
        completion_tokens_details=CompletionTokensDetailsWrapper(
            accepted_prediction_tokens=None,
            audio_tokens=None,
            reasoning_tokens=0,
            rejected_prediction_tokens=None,
            text_tokens=None,
        ),
        prompt_tokens_details=PromptTokensDetailsWrapper(
            audio_tokens=None, cached_tokens=0, text_tokens=None, image_tokens=None
        ),
        cache_creation_input_tokens=118,
        cache_read_input_tokens=28432,
    )

    custom_llm_provider = "vertex_ai"

    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )

    print(f"prompt_cost: {prompt_cost}")
    assert prompt_cost < 0.085


def test_generic_cost_per_token_anthropic_prompt_caching_with_cache_creation():
    model = "claude-haiku-4-5-20251001"
    usage = Usage(
        completion_tokens=90,
        prompt_tokens=28436,
        total_tokens=28526,
        completion_tokens_details=CompletionTokensDetailsWrapper(
            accepted_prediction_tokens=None,
            audio_tokens=None,
            reasoning_tokens=0,
            rejected_prediction_tokens=None,
            text_tokens=None,
        ),
        prompt_tokens_details=None,
        cache_creation_input_tokens=2000,
    )

    custom_llm_provider = "anthropic"

    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )

    print(f"prompt_cost: {prompt_cost}")
    assert round(prompt_cost, 3) == 0.029


def test_string_cost_values():
    """Test that cost values defined as strings are properly converted to floats."""
    from unittest.mock import patch

    # Mock model info with string cost values (as might be read from config.yaml)
    mock_model_info = {
        "input_cost_per_token": "3e-7",  # String representation of scientific notation
        "output_cost_per_token": "6e-7",  # String representation of scientific notation
        "input_cost_per_audio_token": "0.000001",  # String representation of decimal
        "output_cost_per_audio_token": "0.000002",  # String representation of decimal
        "cache_read_input_token_cost": "1.5e-8",  # String representation of scientific notation
        "cache_creation_input_token_cost": "2.5e-8",  # String representation of scientific notation
    }

    # Test usage with various token types
    # Note: prompt_tokens must equal sum of details to avoid double-counting adjustment
    # text_tokens(700) + audio_tokens(100) + cached_tokens(200) + cache_creation_tokens(150) = 1150
    usage = Usage(
        prompt_tokens=1150,
        completion_tokens=500,
        total_tokens=1650,
        prompt_tokens_details=PromptTokensDetailsWrapper(
            audio_tokens=100,
            cached_tokens=200,
            text_tokens=700,
            image_tokens=None,
            cache_creation_tokens=150,
        ),
        completion_tokens_details=CompletionTokensDetailsWrapper(
            audio_tokens=50,
            reasoning_tokens=None,
            text_tokens=450,
            accepted_prediction_tokens=None,
            rejected_prediction_tokens=None,
        ),
    )

    # Mock get_model_info to return our mock model info
    with patch(
        "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
        return_value=mock_model_info,
    ):
        prompt_cost, completion_cost = generic_cost_per_token(
            model="test-model", usage=usage, custom_llm_provider="test-provider"
        )

    # Calculate expected costs manually
    # Prompt cost = text_tokens * input_cost + audio_tokens * audio_cost + cached_tokens * cache_read_cost + cache_creation_tokens * cache_creation_cost
    expected_prompt_cost = (
        700 * 3e-7  # text tokens
        + 100 * 1e-6  # audio tokens
        + 200 * 1.5e-8  # cached tokens
        + 150 * 2.5e-8  # cache creation tokens
    )

    # Completion cost = text_tokens * output_cost + audio_tokens * audio_output_cost
    expected_completion_cost = 450 * 6e-7 + 50 * 2e-6  # text tokens  # audio tokens

    # Assert costs are calculated correctly
    assert round(prompt_cost, 12) == round(expected_prompt_cost, 12)
    assert round(completion_cost, 12) == round(expected_completion_cost, 12)


def test_calculate_cost_component_with_string_values():
    """Test the calculate_cost_component function directly with string cost values."""
    from litellm.litellm_core_utils.llm_cost_calc.utils import calculate_cost_component

    # Test with valid string scientific notation
    model_info = {"input_cost_per_token": "3e-7"}
    cost = calculate_cost_component(model_info, "input_cost_per_token", 1000)
    assert cost == 1000 * 3e-7

    # Test with valid string decimal notation
    model_info = {"output_cost_per_token": "0.000001"}
    cost = calculate_cost_component(model_info, "output_cost_per_token", 500)
    assert cost == 500 * 0.000001

    # Test with float value (should work as before)
    model_info = {"input_cost_per_token": 3e-7}
    cost = calculate_cost_component(model_info, "input_cost_per_token", 1000)
    assert cost == 1000 * 3e-7

    # Test with invalid string value (should return 0.0)
    model_info = {"input_cost_per_token": "invalid_number"}
    cost = calculate_cost_component(model_info, "input_cost_per_token", 1000)
    assert cost == 0.0

    # Test with None value (should return 0.0)
    model_info = {"input_cost_per_token": None}
    cost = calculate_cost_component(model_info, "input_cost_per_token", 1000)
    assert cost == 0.0

    # Test with missing key (should return 0.0)
    model_info = {}
    cost = calculate_cost_component(model_info, "input_cost_per_token", 1000)
    assert cost == 0.0

    # Test with zero usage (should return 0.0)
    model_info = {"input_cost_per_token": "3e-7"}
    cost = calculate_cost_component(model_info, "input_cost_per_token", 0)
    assert cost == 0.0

    # Test with None usage (should return 0.0)
    model_info = {"input_cost_per_token": "3e-7"}
    cost = calculate_cost_component(model_info, "input_cost_per_token", None)
    assert cost == 0.0


def test_string_cost_values_edge_cases():
    """Test edge cases for string cost value handling."""
    from unittest.mock import patch

    # Test with mixed string and float cost values
    mock_model_info = {
        "input_cost_per_token": "1e-6",  # String
        "output_cost_per_token": 2e-6,  # Float
        "input_cost_per_audio_token": "invalid",  # Invalid string
        "output_cost_per_audio_token": None,  # None value
    }

    usage = Usage(
        prompt_tokens=1000,
        completion_tokens=500,
        total_tokens=1500,
        prompt_tokens_details=PromptTokensDetailsWrapper(
            audio_tokens=100, cached_tokens=0, text_tokens=1000, image_tokens=None
        ),
        completion_tokens_details=CompletionTokensDetailsWrapper(
            audio_tokens=50,
            reasoning_tokens=None,
            text_tokens=500,
            accepted_prediction_tokens=None,
            rejected_prediction_tokens=None,
        ),
    )

    with patch(
        "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
        return_value=mock_model_info,
    ):
        prompt_cost, completion_cost = generic_cost_per_token(
            model="test-model", usage=usage, custom_llm_provider="test-provider"
        )

    # Expected costs:
    # Prompt: 1000 * 1e-6 + 100 * 0 (invalid string becomes 0)
    # Completion: 500 * 2e-6 + 50 * 2e-6 (audio tokens fall back to base cost when output_cost_per_audio_token is None)
    expected_prompt_cost = 1000 * 1e-6
    expected_completion_cost = 500 * 2e-6 + 50 * 2e-6

    assert round(prompt_cost, 12) == round(expected_prompt_cost, 12)
    assert round(completion_cost, 12) == round(expected_completion_cost, 12)


def test_string_cost_values_with_threshold():
    """Test that string cost values work correctly with threshold pricing."""
    from unittest.mock import patch

    # Mock model info with string cost values including threshold pricing
    mock_model_info = {
        "input_cost_per_token": "1e-6",  # String base cost
        "output_cost_per_token": "2e-6",  # String base cost
        "input_cost_per_token_above_200k_tokens": "5e-7",  # String threshold cost (lower)
        "output_cost_per_token_above_200k_tokens": "1e-6",  # String threshold cost (lower)
    }

    # Test usage above threshold
    usage = Usage(
        prompt_tokens=250000,  # Above 200k threshold
        completion_tokens=1000,
        total_tokens=251000,
    )

    with patch(
        "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info",
        return_value=mock_model_info,
    ):
        prompt_cost, completion_cost = generic_cost_per_token(
            model="test-model", usage=usage, custom_llm_provider="test-provider"
        )

    # Expected costs using threshold pricing (string values converted to float)
    expected_prompt_cost = 250000 * 5e-7  # threshold cost
    expected_completion_cost = 1000 * 1e-6  # threshold cost

    assert round(prompt_cost, 12) == round(expected_prompt_cost, 12)
    assert round(completion_cost, 12) == round(expected_completion_cost, 12)


def test_calculate_cache_writing_cost():
    """Test the calculate_cache_writing_cost function with detailed cache creation token breakdown."""

    # Test case 1: With cache creation token details (matching the provided input)
    cache_creation_tokens = 14055
    cache_creation_token_details = CacheCreationTokenDetails(
        ephemeral_5m_input_tokens=56, ephemeral_1h_input_tokens=13999
    )
    cache_creation_cost_above_1hr = 6e-06
    cache_creation_cost = 3.75e-06

    result = calculate_cache_writing_cost(
        cache_creation_tokens=cache_creation_tokens,
        cache_creation_token_details=cache_creation_token_details,
        cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
        cache_creation_cost=cache_creation_cost,
    )

    # Expected calculation:
    # 5m tokens: 56 * 3.75e-06 = 0.00021
    # 1h tokens: 13999 * 6e-06 = 0.083994
    # Total: 0.00021 + 0.083994 = 0.084204
    expected_cost = (56 * 3.75e-06) + (13999 * 6e-06)

    assert round(result, 6) == round(expected_cost, 6)
    assert round(result, 6) == 0.084204

    # Test case 2: Without cache creation token details (fallback behavior)
    cache_creation_tokens_no_details = 1000
    cache_creation_token_details_none = None
    cache_creation_cost_fallback = 5e-06

    result_no_details = calculate_cache_writing_cost(
        cache_creation_tokens=cache_creation_tokens_no_details,
        cache_creation_token_details=cache_creation_token_details_none,
        cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
        cache_creation_cost=cache_creation_cost_fallback,
    )

    # Expected calculation when no details: 1000 * 5e-06 = 0.005
    expected_cost_no_details = 1000 * 5e-06

    assert round(result_no_details, 6) == round(expected_cost_no_details, 6)
    assert result_no_details == 0.005

    # Test case 3: With cache creation token details but None values
    cache_creation_token_details_partial = CacheCreationTokenDetails(
        ephemeral_5m_input_tokens=None, ephemeral_1h_input_tokens=100
    )

    result_partial = calculate_cache_writing_cost(
        cache_creation_tokens=500,
        cache_creation_token_details=cache_creation_token_details_partial,
        cache_creation_cost_above_1hr=6e-06,
        cache_creation_cost=3e-06,
    )

    # Expected calculation: 0 (for None 5m tokens) + (100 * 6e-06) = 0.0006
    expected_cost_partial = (0.0) + (100 * 6e-06)

    assert round(result_partial, 6) == round(expected_cost_partial, 6)
    assert round(result_partial, 6) == 0.0006

    # Test case 4: Zero costs
    result_zero = calculate_cache_writing_cost(
        cache_creation_tokens=1000,
        cache_creation_token_details=CacheCreationTokenDetails(
            ephemeral_5m_input_tokens=50, ephemeral_1h_input_tokens=950
        ),
        cache_creation_cost_above_1hr=0.0,
        cache_creation_cost=0.0,
    )

    assert result_zero == 0.0


def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details():
    """
    Regression test: when cache_creation_tokens is 0 but cache_creation_token_details
    has non-zero ephemeral tokens, the cost must still be calculated.
    This ensures the guard in _calculate_input_cost doesn't skip
    calculate_cache_writing_cost when only ephemeral token details are present.
    """
    cache_creation_cost = 3.75e-06
    cache_creation_cost_above_1hr = 6e-06

    prompt_tokens_details: PromptTokensDetailsResult = {
        "cache_hit_tokens": 0,
        "cache_creation_tokens": 0,
        "cache_creation_token_details": CacheCreationTokenDetails(
            ephemeral_5m_input_tokens=100,
            ephemeral_1h_input_tokens=200,
        ),
        "text_tokens": 0,
        "audio_tokens": 0,
        "image_tokens": 0,
        "character_count": 0,
        "image_count": 0,
        "video_length_seconds": 0.0,
    }

    model_info: ModelInfo = {}

    result = _calculate_input_cost(
        prompt_tokens_details=prompt_tokens_details,
        model_info=model_info,
        prompt_base_cost=0.0,
        cache_read_cost=0.0,
        cache_creation_cost=cache_creation_cost,
        cache_creation_cost_above_1hr=cache_creation_cost_above_1hr,
    )

    # Expected: (100 * 3.75e-06) + (200 * 6e-06) = 0.000375 + 0.0012 = 0.001575
    expected = (100 * cache_creation_cost) + (200 * cache_creation_cost_above_1hr)
    assert (
        result > 0
    ), "Cost should not be zero when ephemeral token details are present"
    assert round(result, 6) == round(expected, 6)


def test_service_tier_flex_pricing():
    """Test that flex service tier uses correct pricing (approximately 50% of standard)."""
    # Set up environment for local model cost map
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    # Test with gpt-5-nano which has flex pricing
    model = "gpt-5-nano"
    custom_llm_provider = "openai"

    # Create usage object
    usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)

    # Test standard pricing
    std_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
        service_tier=None,
    )
    std_total = std_cost[0] + std_cost[1]

    # Test flex pricing
    flex_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
        service_tier="flex",
    )
    flex_total = flex_cost[0] + flex_cost[1]

    # Verify flex is approximately 50% of standard
    assert std_total > 0, "Standard cost should be greater than 0"
    assert flex_total > 0, "Flex cost should be greater than 0"

    flex_ratio = flex_total / std_total
    assert (
        0.45 <= flex_ratio <= 0.55
    ), f"Flex pricing should be ~50% of standard, got {flex_ratio:.2f}"

    # Verify specific costs match expected values
    # gpt-5-nano flex: input=2.5e-08, output=2e-07
    expected_flex_prompt = 1000 * 2.5e-08  # 0.000025
    expected_flex_completion = 500 * 2e-07  # 0.0001
    expected_flex_total = expected_flex_prompt + expected_flex_completion

    assert (
        abs(flex_cost[0] - expected_flex_prompt) < 1e-10
    ), f"Flex prompt cost mismatch: {flex_cost[0]} vs {expected_flex_prompt}"
    assert (
        abs(flex_cost[1] - expected_flex_completion) < 1e-10
    ), f"Flex completion cost mismatch: {flex_cost[1]} vs {expected_flex_completion}"
    assert (
        abs(flex_total - expected_flex_total) < 1e-10
    ), f"Flex total cost mismatch: {flex_total} vs {expected_flex_total}"


def test_service_tier_default_pricing():
    """Test that when no service tier is provided, standard pricing is used."""
    # Set up environment for local model cost map
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    # Test with gpt-5-nano
    model = "gpt-5-nano"
    custom_llm_provider = "openai"

    # Create usage object
    usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)

    # Test with no service tier (should use standard)
    default_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
        service_tier=None,
    )

    # Test with explicit standard service tier
    standard_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
        service_tier="standard",
    )

    # Both should be identical
    assert (
        abs(default_cost[0] - standard_cost[0]) < 1e-10
    ), "Default and standard prompt costs should be identical"
    assert (
        abs(default_cost[1] - standard_cost[1]) < 1e-10
    ), "Default and standard completion costs should be identical"

    # Verify specific costs match expected standard values
    # gpt-5-nano standard: input=5e-08, output=4e-07
    expected_standard_prompt = 1000 * 5e-08  # 0.00005
    expected_standard_completion = 500 * 4e-07  # 0.0002
    expected_standard_total = expected_standard_prompt + expected_standard_completion

    assert (
        abs(default_cost[0] - expected_standard_prompt) < 1e-10
    ), f"Standard prompt cost mismatch: {default_cost[0]} vs {expected_standard_prompt}"
    assert (
        abs(default_cost[1] - expected_standard_completion) < 1e-10
    ), f"Standard completion cost mismatch: {default_cost[1]} vs {expected_standard_completion}"


def test_service_tier_fallback_pricing():
    """Test that when service tier is provided but model doesn't have those keys, it falls back to standard pricing."""
    # Set up environment for local model cost map
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    # Test with gpt-4 which doesn't have flex pricing keys
    model = "gpt-4"
    custom_llm_provider = "openai"

    # Create usage object
    usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)

    # Test standard pricing
    std_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
        service_tier=None,
    )
    std_total = std_cost[0] + std_cost[1]

    # Test flex pricing (should fall back to standard since gpt-4 doesn't have flex keys)
    flex_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
        service_tier="flex",
    )
    flex_total = flex_cost[0] + flex_cost[1]

    # Test priority pricing (should fall back to standard since gpt-4 doesn't have priority keys)
    priority_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
        service_tier="priority",
    )
    priority_total = priority_cost[0] + priority_cost[1]

    # All should be identical (fallback to standard)
    assert (
        abs(std_total - flex_total) < 1e-10
    ), f"Standard and flex costs should be identical (fallback): {std_total} vs {flex_total}"
    assert (
        abs(std_total - priority_total) < 1e-10
    ), f"Standard and priority costs should be identical (fallback): {std_total} vs {priority_total}"

    # Verify costs are reasonable (not zero)
    assert std_total > 0, "Standard cost should be greater than 0"
    assert flex_total > 0, "Flex cost should be greater than 0 (fallback)"
    assert priority_total > 0, "Priority cost should be greater than 0 (fallback)"

    # Verify specific costs match expected gpt-4 values
    # gpt-4 standard: input=3e-05, output=6e-05
    expected_standard_prompt = 1000 * 3e-05  # 0.03
    expected_standard_completion = 500 * 6e-05  # 0.03
    expected_standard_total = expected_standard_prompt + expected_standard_completion

    assert (
        abs(std_cost[0] - expected_standard_prompt) < 1e-10
    ), f"Standard prompt cost mismatch: {std_cost[0]} vs {expected_standard_prompt}"
    assert (
        abs(std_cost[1] - expected_standard_completion) < 1e-10
    ), f"Standard completion cost mismatch: {std_cost[1]} vs {expected_standard_completion}"


@pytest.mark.parametrize(
    "model",
    [
        "gemini-3-pro-image-preview",
        "gemini-3.1-flash-image-preview",
    ],
)
def test_gemini_image_generation_cost_with_zero_text_tokens(model: str):
    """
    Test that image_tokens are correctly costed when text_tokens=0.

    Reproduces issue #17410: completion_cost calculates incorrectly for
    Gemini-3-pro-image model - image_tokens were treated as text tokens
    when text_tokens=0.

    https://github.com/BerriAI/litellm/issues/17410
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    custom_llm_provider = "vertex_ai"

    # Usage from the issue: text_tokens=0, image_tokens=1120, reasoning_tokens=225
    usage = Usage(
        completion_tokens=1345,
        prompt_tokens=10,
        total_tokens=1355,
        completion_tokens_details=CompletionTokensDetailsWrapper(
            accepted_prediction_tokens=None,
            audio_tokens=None,
            reasoning_tokens=225,
            rejected_prediction_tokens=None,
            text_tokens=0,  # This is the key: text_tokens=0
            image_tokens=1120,
        ),
        prompt_tokens_details=PromptTokensDetailsWrapper(
            audio_tokens=None, cached_tokens=None, text_tokens=10, image_tokens=None
        ),
    )

    model_cost_map = litellm.model_cost[model]
    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )

    # Expected costs:
    # - text_tokens: 0 * output_cost_per_token = 0
    # - image_tokens: 1120 * output_cost_per_image_token
    # - reasoning_tokens: 225 * output_cost_per_token
    # Total completion should include both image + reasoning costs.

    output_cost_per_image_token = model_cost_map.get("output_cost_per_image_token", 0)
    output_cost_per_token = model_cost_map.get("output_cost_per_token", 0)

    expected_image_cost = 1120 * output_cost_per_image_token
    expected_reasoning_cost = (
        225 * output_cost_per_token
    )  # reasoning uses base token cost
    expected_completion_cost = expected_image_cost + expected_reasoning_cost

    # The bug was: all completion tokens were treated as text tokens only.
    bugged_text_only_cost = 1345 * output_cost_per_token
    assert completion_cost > bugged_text_only_cost * 2, (
        f"Completion cost should be significantly larger than text-only bugged path. "
        f"Expected > {bugged_text_only_cost * 2:.6f}, got {completion_cost:.6f}"
    )
    assert round(completion_cost, 4) == round(
        expected_completion_cost, 4
    ), f"Expected completion cost ${expected_completion_cost:.6f}, got ${completion_cost:.6f}"


def test_vertex_image_generation_cost_prefers_token_usage_metadata():
    """
    When usage metadata exists on image responses, Vertex image generation cost
    should be calculated from token pricing, not flat output_cost_per_image.
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    model = "gemini-3.1-flash-image-preview"
    model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")

    input_text_tokens = 50
    input_image_tokens = 1120
    output_image_tokens = 1120
    prompt_tokens = input_text_tokens + input_image_tokens

    image_response = ImageResponse(
        data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")],
        usage=ImageUsage(
            input_tokens=prompt_tokens,
            input_tokens_details=ImageUsageInputTokensDetails(
                text_tokens=input_text_tokens,
                image_tokens=input_image_tokens,
            ),
            output_tokens=output_image_tokens,
            total_tokens=prompt_tokens + output_image_tokens,
        ),
    )

    cost = vertex_image_generation_cost_calculator(
        model=model,
        image_response=image_response,
    )

    expected_prompt_cost = prompt_tokens * model_info["input_cost_per_token"]
    expected_completion_cost = (
        output_image_tokens * model_info["output_cost_per_image_token"]
    )
    expected_total_cost = expected_prompt_cost + expected_completion_cost

    assert round(cost, 10) == round(expected_total_cost, 10)
    # Ensure this is not falling back to flat per-image pricing.
    assert cost != len(image_response.data) * model_info["output_cost_per_image"]


def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing():
    """
    Without usage metadata, Vertex image generation cost should fall back to
    output_cost_per_image * number_of_images.
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    model = "gemini-3.1-flash-image-preview"
    model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")

    image_response = ImageResponse(
        data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]
    )

    cost = vertex_image_generation_cost_calculator(
        model=model,
        image_response=image_response,
    )

    expected_cost = len(image_response.data) * model_info["output_cost_per_image"]
    assert round(cost, 10) == round(expected_cost, 10)


def test_gemini_image_generation_cost_prefers_token_usage_metadata():
    """
    When usage metadata exists on image responses, Gemini image generation cost
    should be calculated from token pricing, not flat output_cost_per_image.
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    model = "gemini/gemini-3-pro-image-preview"
    model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini")

    input_text_tokens = 20
    input_image_tokens = 1120
    output_image_tokens = 1120
    prompt_tokens = input_text_tokens + input_image_tokens

    image_response = ImageResponse(
        data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")],
        usage=ImageUsage(
            input_tokens=prompt_tokens,
            input_tokens_details=ImageUsageInputTokensDetails(
                text_tokens=input_text_tokens,
                image_tokens=input_image_tokens,
            ),
            output_tokens=output_image_tokens,
            total_tokens=prompt_tokens + output_image_tokens,
        ),
    )

    cost = gemini_image_generation_cost_calculator(
        model=model,
        image_response=image_response,
    )

    expected_prompt_cost = prompt_tokens * model_info["input_cost_per_token"]
    expected_completion_cost = (
        output_image_tokens * model_info["output_cost_per_image_token"]
    )
    expected_total_cost = expected_prompt_cost + expected_completion_cost

    assert round(cost, 10) == round(expected_total_cost, 10)
    # Ensure this is not falling back to flat per-image pricing.
    assert cost != len(image_response.data) * model_info["output_cost_per_image"]


def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing():
    """
    Without usage metadata, Gemini image generation cost should fall back to
    output_cost_per_image * number_of_images.
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    model = "gemini/gemini-3-pro-image-preview"
    model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini")

    image_response = ImageResponse(
        data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]
    )

    cost = gemini_image_generation_cost_calculator(
        model=model,
        image_response=image_response,
    )

    expected_cost = len(image_response.data) * model_info["output_cost_per_image"]
    assert round(cost, 10) == round(expected_cost, 10)


def test_bedrock_anthropic_prompt_caching():
    """Test Bedrock Anthropic models with prompt caching return correct costs."""
    model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
    usage = Usage(
        prompt_tokens=52123,
        completion_tokens=497,
        total_tokens=52620,
        cache_creation_input_tokens=7183,
        cache_read_input_tokens=22465,
    )

    custom_llm_provider = "bedrock"

    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )

    assert prompt_cost >= 0
    assert completion_cost >= 0
    assert round(prompt_cost, 3) == 0.111
    assert round(completion_cost, 5) == 0.00820


def test_reasoning_tokens_without_text_tokens_gpt5_nano():
    """
    Test fix for GitHub issue #18599:
    https://github.com/BerriAI/litellm/issues/18599

    When OpenAI models (gpt-5-nano, o1, o3) return reasoning_tokens but don't provide
    text_tokens, LiteLLM should calculate text_tokens as:
      text_tokens = completion_tokens - reasoning_tokens - audio_tokens - image_tokens

    This ensures ALL completion tokens are billed, not just reasoning tokens.
    """
    model = "gpt-5-nano"
    custom_llm_provider = "openai"

    # Simulate OpenAI gpt-5-nano response where text_tokens is NOT provided
    # completion_tokens: 977 total
    # reasoning_tokens: 768
    # text_tokens: should be calculated as 977 - 768 = 209
    usage = Usage(
        prompt_tokens=17,
        completion_tokens=977,
        total_tokens=994,
        completion_tokens_details=CompletionTokensDetailsWrapper(
            reasoning_tokens=768,
            audio_tokens=0,
            # text_tokens NOT provided - this is the key part of the bug
        ),
    )

    prompt_cost, completion_cost = generic_cost_per_token(
        model=model,
        usage=usage,
        custom_llm_provider=custom_llm_provider,
    )

    # gpt-5-nano pricing: $0.05/1M input, $0.40/1M output
    expected_prompt_cost = 17 * 0.05 / 1_000_000
    expected_completion_cost = 977 * 0.40 / 1_000_000  # ALL tokens, not just reasoning

    assert (
        abs(prompt_cost - expected_prompt_cost) < 1e-10
    ), f"Prompt cost incorrect: {prompt_cost} vs {expected_prompt_cost}"

    assert (
        abs(completion_cost - expected_completion_cost) < 1e-10
    ), f"Completion cost incorrect: {completion_cost} vs {expected_completion_cost}"

    # Verify it's NOT using only reasoning_tokens (the bug)
    wrong_cost = 768 * 0.40 / 1_000_000  # Only reasoning tokens
    assert (
        abs(completion_cost - wrong_cost) > 1e-6
    ), "Bug detected: Cost calculation is using only reasoning_tokens instead of all completion_tokens!"


def test_image_count_prevents_text_tokens_fallback():
    """
    Test that the text_tokens fallback in generic_cost_per_token does not
    override text_tokens=0 when image_count > 0.

    Regression test for: Bedrock image embedding double-charging bug.
    When image_count > 0, text_tokens=0 is intentional (image-only request),
    not "text_tokens not set by provider."
    """
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")

    # Simulate Nova image-only embedding: prompt_tokens estimated from
    # embedding dimensions (768 for 3072-dim), image_count=1
    usage = Usage(
        prompt_tokens=768,
        completion_tokens=0,
        total_tokens=768,
        prompt_tokens_details=PromptTokensDetailsWrapper(
            image_count=1,
        ),
    )

    prompt_cost, completion_cost = generic_cost_per_token(
        model="amazon.nova-2-multimodal-embeddings-v1:0",
        usage=usage,
        custom_llm_provider="bedrock",
    )

    # Cost should be 1 * input_cost_per_image ($6e-05) = $0.00006
    # NOT 768 * input_cost_per_token ($1.35e-07) + $0.00006 = $0.000164
    expected_image_cost = 1 * 6e-05
    assert prompt_cost == expected_image_cost, (
        f"Expected prompt_cost={expected_image_cost} (image-only), "
        f"got {prompt_cost}. text_tokens fallback may be double-charging."
    )
    assert completion_cost == 0.0


# ---------------------------------------------------------------------------
# Data-residency (OpenAI regional processing) tests
# ---------------------------------------------------------------------------


@pytest.fixture
def _local_model_cost_map():
    prev_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP")
    prev_model_cost = litellm.model_cost
    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
    litellm.model_cost = litellm.get_model_cost_map(url="")
    try:
        yield
    finally:
        litellm.model_cost = prev_model_cost
        if prev_env is None:
            os.environ.pop("LITELLM_LOCAL_MODEL_COST_MAP", None)
        else:
            os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = prev_env


@pytest.mark.parametrize("data_residency", ["eu", "us"])
def test_data_residency_applies_uplift(data_residency, _local_model_cost_map):
    """gpt-5 should apply the regional processing uplift multiplier when
    data_residency is set."""
    from litellm.types.utils import Usage

    usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)

    base = generic_cost_per_token(
        model="gpt-5",
        usage=usage,
        custom_llm_provider="openai",
    )
    regional = generic_cost_per_token(
        model="gpt-5",
        usage=usage,
        custom_llm_provider="openai",
        data_residency=data_residency,
    )

    base_total = base[0] + base[1]
    regional_total = regional[0] + regional[1]

    assert base_total > 0
    assert regional_total == pytest.approx(base_total * 1.10, rel=1e-9)
    assert regional[0] == pytest.approx(base[0] * 1.10, rel=1e-9)
    assert regional[1] == pytest.approx(base[1] * 1.10, rel=1e-9)


def test_data_residency_no_uplift_for_unmarked_model(_local_model_cost_map):
    """A model without a regional_processing_uplift_multiplier_* entry should
    fall back to base pricing, not error."""
    from litellm.types.utils import Usage

    usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)

    base = generic_cost_per_token(
        model="gpt-3.5-turbo",
        usage=usage,
        custom_llm_provider="openai",
    )
    with_residency = generic_cost_per_token(
        model="gpt-3.5-turbo",
        usage=usage,
        custom_llm_provider="openai",
        data_residency="eu",
    )

    assert base == with_residency


def test_data_residency_none_no_uplift(_local_model_cost_map):
    """data_residency=None should be a no-op even for models with a multiplier."""
    from litellm.types.utils import Usage

    usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)

    base = generic_cost_per_token(
        model="gpt-5",
        usage=usage,
        custom_llm_provider="openai",
    )
    explicit_none = generic_cost_per_token(
        model="gpt-5",
        usage=usage,
        custom_llm_provider="openai",
        data_residency=None,
    )

    assert base == explicit_none


def test_data_residency_composes_with_service_tier(_local_model_cost_map):
    """The uplift multiplies the priority-tier cost, not the standard one."""
    from litellm.types.utils import Usage

    usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)

    priority_base = generic_cost_per_token(
        model="gpt-5",
        usage=usage,
        custom_llm_provider="openai",
        service_tier="priority",
    )
    priority_eu = generic_cost_per_token(
        model="gpt-5",
        usage=usage,
        custom_llm_provider="openai",
        service_tier="priority",
        data_residency="eu",
    )

    priority_base_total = priority_base[0] + priority_base[1]
    priority_eu_total = priority_eu[0] + priority_eu[1]

    assert priority_base_total > 0
    assert priority_eu_total == pytest.approx(priority_base_total * 1.10, rel=1e-9)
