"""
Tests for OpenAI gpt-image-1 cost calculator

This tests the fix for GitHub issue #13847:
https://github.com/BerriAI/litellm/issues/13847

gpt-image-1 uses token-based pricing:
- Text Input: $5.00/1M tokens
- Image Input: $10.00/1M tokens
- Image Output: $40.00/1M tokens
"""



import pytest

import litellm
from litellm.types.utils import (
    CompletionTokensDetailsWrapper,
    ImageResponse,
    ImageObject,
    ImageUsage,
    ImageUsageInputTokensDetails,
    PromptTokensDetailsWrapper,
    Usage,
)


@pytest.fixture(autouse=True)
def _use_local_model_cost_map(monkeypatch):
    original_model_cost = litellm.model_cost
    monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
    litellm.model_cost = litellm.get_model_cost_map(url="")
    litellm.get_model_info.cache_clear()
    try:
        yield
    finally:
        litellm.model_cost = original_model_cost
        litellm.get_model_info.cache_clear()


class TestGPTImageCostCalculator:
    """Test the OpenAI gpt-image cost calculator"""

    def test_gpt_image_1_cost_with_text_only(self):
        """Test cost calculation with only text input tokens"""
        from litellm.llms.openai.image_generation.cost_calculator import cost_calculator

        usage = ImageUsage(
            input_tokens=100,
            output_tokens=5000,
            total_tokens=5100,
            input_tokens_details=ImageUsageInputTokensDetails(
                text_tokens=100,
                image_tokens=0,
            ),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(url="http://example.com/image.jpg")],
        )
        image_response.usage = usage

        cost = cost_calculator(
            model="gpt-image-1",
            image_response=image_response,
            custom_llm_provider="openai",
        )

        # Expected cost:
        # Text input: 100 * $5/1M = 0.0005
        # Image output: 5000 * $40/1M = 0.2
        # Total: 0.2005
        expected_cost = 0.0005 + 0.2
        assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"

    def test_gpt_image_1_cost_with_image_input(self):
        """Test cost calculation with both text and image input tokens (for edits)"""
        from litellm.llms.openai.image_generation.cost_calculator import cost_calculator

        usage = ImageUsage(
            input_tokens=600,
            output_tokens=5000,
            total_tokens=5600,
            input_tokens_details=ImageUsageInputTokensDetails(
                text_tokens=100,
                image_tokens=500,
            ),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(url="http://example.com/image.jpg")],
        )
        image_response.usage = usage

        cost = cost_calculator(
            model="gpt-image-1",
            image_response=image_response,
            custom_llm_provider="openai",
        )

        # Expected cost:
        # Text input: 100 * $5/1M = 0.0005
        # Image input: 500 * $10/1M = 0.005
        # Image output: 5000 * $40/1M = 0.2
        # Total: 0.2055
        expected_cost = 0.0005 + 0.005 + 0.2
        assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"

    def test_gpt_image_1_mini_cost(self):
        """Test cost calculation for gpt-image-1-mini model"""
        from litellm.llms.openai.image_generation.cost_calculator import cost_calculator

        usage = ImageUsage(
            input_tokens=100,
            output_tokens=5000,
            total_tokens=5100,
            input_tokens_details=ImageUsageInputTokensDetails(
                text_tokens=100,
                image_tokens=0,
            ),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(url="http://example.com/image.jpg")],
        )
        image_response.usage = usage

        cost = cost_calculator(
            model="gpt-image-1-mini",
            image_response=image_response,
            custom_llm_provider="openai",
        )

        # Expected cost for gpt-image-1-mini:
        # Text input: 100 * $2/1M = 0.0002
        # Image output: 5000 * $8/1M = 0.04
        # Total: 0.0402
        expected_cost = 0.0002 + 0.04
        assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"

    def test_gpt_image_1_cost_no_usage(self):
        """Test that cost returns 0 when no usage data is available"""
        from litellm.llms.openai.image_generation.cost_calculator import cost_calculator

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(url="http://example.com/image.jpg")],
        )

        cost = cost_calculator(
            model="gpt-image-1",
            image_response=image_response,
            custom_llm_provider="openai",
        )

        assert cost == 0.0

    def test_gpt_image_2_cost_with_text_and_image_tokens(self):
        """Test cost calculation for gpt-image-2 token pricing"""
        from litellm.llms.openai.image_generation.cost_calculator import cost_calculator

        usage = Usage(
            prompt_tokens=600,
            completion_tokens=5000,
            total_tokens=5600,
            prompt_tokens_details=PromptTokensDetailsWrapper(
                text_tokens=100,
                image_tokens=500,
            ),
            completion_tokens_details=CompletionTokensDetailsWrapper(
                text_tokens=1000,
                image_tokens=4000,
            ),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(url="http://example.com/image.jpg")],
        )
        image_response.usage = usage

        cost = cost_calculator(
            model="gpt-image-2",
            image_response=image_response,
            custom_llm_provider="openai",
        )

        # GPT Image 2 pricing:
        # Text input: 100 * $5/1M = 0.0005
        # Image input: 500 * $8/1M = 0.004
        # Text output: 1000 * $10/1M = 0.01
        # Image output: 4000 * $30/1M = 0.12
        expected_cost = 0.0005 + 0.004 + 0.01 + 0.12
        assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"


class TestGPTImageCostRouting:
    """Test that gpt-image models are properly routed to the token-based calculator"""

    def test_openai_gpt_image_routes_to_token_calculator(self):
        """Test that OpenAI gpt-image-1 routes to token-based calculator"""
        from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils

        usage = ImageUsage(
            input_tokens=100,
            output_tokens=5000,
            total_tokens=5100,
            input_tokens_details=ImageUsageInputTokensDetails(
                text_tokens=100,
                image_tokens=0,
            ),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(url="http://example.com/image.jpg")],
        )
        image_response.usage = usage

        cost = CostCalculatorUtils.route_image_generation_cost_calculator(
            model="gpt-image-1",
            completion_response=image_response,
            custom_llm_provider="openai",
        )

        expected_cost = 0.0005 + 0.2
        assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"

    def test_openai_gpt_image_2_routes_to_token_calculator(self):
        """Test that OpenAI gpt-image-2 routes to token-based calculator"""
        from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils

        usage = Usage(
            prompt_tokens=100,
            completion_tokens=5000,
            total_tokens=5100,
            prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=100),
            completion_tokens_details=CompletionTokensDetailsWrapper(image_tokens=5000),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(url="http://example.com/image.jpg")],
        )
        image_response.usage = usage

        cost = CostCalculatorUtils.route_image_generation_cost_calculator(
            model="gpt-image-2",
            completion_response=image_response,
            custom_llm_provider="openai",
        )

        expected_cost = 0.0005 + 0.15
        assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"

    def test_openai_dalle_routes_to_pixel_calculator(self):
        """Test that OpenAI DALL-E still routes to pixel-based calculator"""
        from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(url="http://example.com/image.jpg")],
        )
        image_response.size = "1024x1024"
        image_response.quality = "standard"

        cost = CostCalculatorUtils.route_image_generation_cost_calculator(
            model="dall-e-3",
            completion_response=image_response,
            custom_llm_provider="openai",
            size="1024x1024",
            quality="standard",
            n=1,
        )

        assert cost >= 0


class TestGPTImage15OutputImageTokens:
    """
    Test for GitHub issue #19508:
    Image usage calculation does not include image tokens in gpt-image-1.5

    gpt-image-1.5 returns output_tokens_details with separate image_tokens and text_tokens,
    and these must be correctly included in cost calculation.
    """

    def test_gpt_image_15_output_image_tokens_cost(self):
        """
        Test that output image tokens are correctly included in cost calculation.

        This tests the fix for issue #19508 where output_tokens_details.image_tokens
        were not being included in the cost calculation, causing costs to be
        underreported (e.g., $0.046 instead of $0.14).
        """
        # Simulate gpt-image-1.5 response with output_tokens_details
        # This is what the API returns and what convert_to_image_response transforms
        usage = Usage(
            prompt_tokens=169,
            completion_tokens=4599,
            total_tokens=4768,
            prompt_tokens_details=PromptTokensDetailsWrapper(
                text_tokens=169,
                image_tokens=0,
            ),
            completion_tokens_details=CompletionTokensDetailsWrapper(
                text_tokens=439,
                image_tokens=4160,
            ),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(b64_json="test")],
        )
        image_response.usage = usage
        image_response._hidden_params = {"custom_llm_provider": "openai"}

        cost = litellm.completion_cost(
            completion_response=image_response,
            model="gpt-image-1.5",
            call_type="image_generation",
            custom_llm_provider="openai",
        )

        # gpt-image-1.5 pricing:
        # - input_cost_per_token: 5e-06 ($5/1M for text input)
        # - output_cost_per_token: 1e-05 ($10/1M for text output)
        # - output_cost_per_image_token: 3.2e-05 ($32/1M for image output)
        #
        # Expected cost:
        # Input text: 169 * $5/1M = $0.000845
        # Output text: 439 * $10/1M = $0.00439
        # Output image: 4160 * $32/1M = $0.13312
        # Total: $0.138355
        expected_cost = 169 * 5e-06 + 439 * 1e-05 + 4160 * 3.2e-05

        assert abs(cost - expected_cost) < 1e-6, (
            f"Expected {expected_cost}, got {cost}. "
            f"Image tokens may not be included in cost calculation."
        )


class TestCompletionCostIntegration:
    """Test the full completion_cost integration for gpt-image-1"""

    def test_completion_cost_gpt_image_1(self):
        """Test completion_cost correctly calculates gpt-image-1 costs"""
        usage = ImageUsage(
            input_tokens=100,
            output_tokens=5000,
            total_tokens=5100,
            input_tokens_details=ImageUsageInputTokensDetails(
                text_tokens=100,
                image_tokens=0,
            ),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(url="http://example.com/image.jpg")],
        )
        image_response.usage = usage
        image_response._hidden_params = {"custom_llm_provider": "openai"}

        cost = litellm.completion_cost(
            completion_response=image_response,
            model="gpt-image-1",
            call_type="image_generation",
            custom_llm_provider="openai",
        )

        expected_cost = 0.0005 + 0.2
        assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"


class TestGPTImage2OutputImageTokensNoBreakdown:
    """
    Regression test: the OpenAI Images endpoints (/v1/images/generations and
    /v1/images/edits) return usage with NO output token breakdown — litellm's
    ImageUsage has no ``output_tokens_details`` field. Before the fix, the
    generated-image OUTPUT tokens were priced at the text rate
    (``output_cost_per_token`` = $10/1M for gpt-image-2) instead of the image rate
    (``output_cost_per_image_token`` = $30/1M), a ~3x undercount on the dominant
    cost component.
    """

    def test_gpt_image_2_output_priced_as_image_when_no_breakdown(self):
        from litellm.llms.openai.image_generation.cost_calculator import (
            cost_calculator,
        )

        # Mirrors a real gpt-image-2 /v1/images/edits response: input breakdown is
        # present, but there is no usable output token breakdown.
        usage = ImageUsage(
            input_tokens=3987,
            output_tokens=5488,
            total_tokens=9475,
            input_tokens_details=ImageUsageInputTokensDetails(
                text_tokens=943,
                image_tokens=3044,
            ),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(b64_json="test")],
        )
        image_response.usage = usage
        image_response._hidden_params = {"custom_llm_provider": "openai"}

        cost = cost_calculator(
            model="gpt-image-2",
            image_response=image_response,
            custom_llm_provider="openai",
        )

        # gpt-image-2 pricing:
        #   text input:   943  * $5/1M  = 0.004715
        #   image input:  3044 * $8/1M  = 0.024352
        #   image output: 5488 * $30/1M = 0.164640  (NOT text output $10/1M = 0.054880)
        expected_cost = 943 * 5e-6 + 3044 * 8e-6 + 5488 * 3e-5
        assert abs(cost - expected_cost) < 1e-6, (
            f"Expected {expected_cost}, got {cost}. Generated image output tokens "
            f"are likely being priced at the text output_cost_per_token rate."
        )

    def test_gpt_image_2_chat_usage_without_breakdown_is_costed_not_zero(self):
        """A chat ``Usage`` with ``completion_tokens_details=None`` must still be
        costed via ``generic_cost_per_token`` (output at the text rate) rather than
        erroring or silently returning 0.0."""
        from litellm.llms.openai.image_generation.cost_calculator import (
            cost_calculator,
        )

        usage = Usage(
            prompt_tokens=600,
            completion_tokens=5000,
            total_tokens=5600,
            prompt_tokens_details=PromptTokensDetailsWrapper(
                text_tokens=100,
                image_tokens=500,
            ),
        )

        image_response = ImageResponse(
            created=1234567890,
            data=[ImageObject(b64_json="test")],
        )
        image_response.usage = usage
        image_response._hidden_params = {"custom_llm_provider": "openai"}

        cost = cost_calculator(
            model="gpt-image-2",
            image_response=image_response,
            custom_llm_provider="openai",
        )

        # No output breakdown -> output priced at the text rate (output_cost_per_token):
        #   text in 100*$5/1M + image in 500*$8/1M + output 5000*$10/1M
        expected_cost = 100 * 5e-6 + 500 * 8e-6 + 5000 * 1e-5
        assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
