import os
import sys
import json
import pytest

sys.path.insert(
    0, os.path.abspath("../..")
)  # Adds the parent directory to the system path
import litellm
from litellm import transcription
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIConfig
from base_llm_unit_tests import BaseLLMChatTest
from base_audio_transcription_unit_tests import BaseLLMAudioTranscriptionTest

fireworks = FireworksAIConfig()


def test_map_openai_params_tool_choice():
    # Test case 1: tool_choice is "required"
    result = fireworks.map_openai_params(
        {"tool_choice": "required"}, {}, "some_model", drop_params=False
    )
    assert result == {"tool_choice": "any"}

    # Test case 2: tool_choice is "auto"
    result = fireworks.map_openai_params(
        {"tool_choice": "auto"}, {}, "some_model", drop_params=False
    )
    assert result == {"tool_choice": "auto"}

    # Test case 3: tool_choice is not present
    result = fireworks.map_openai_params(
        {"some_other_param": "value"}, {}, "some_model", drop_params=False
    )
    assert result == {}

    # Test case 4: tool_choice is None
    result = fireworks.map_openai_params(
        {"tool_choice": None}, {}, "some_model", drop_params=False
    )
    assert result == {"tool_choice": None}


def test_map_response_format():
    """
    json_schema response_format is passed through to Fireworks unchanged.

    Fireworks accepts the OpenAI strict json_schema shape natively. The earlier
    downgrade to {type: json_object, schema: ...} silently dropped `strict` and
    `name`, producing a request that Fireworks treats as "any valid JSON" per
    its docs, disabling grammar-guided decoding.

    Ref: https://docs.fireworks.ai/structured-responses/structured-response-formatting
    """
    response_format = {
        "type": "json_schema",
        "json_schema": {
            "schema": {
                "properties": {"result": {"type": "boolean"}},
                "required": ["result"],
                "type": "object",
            },
            "name": "BooleanResponse",
            "strict": True,
        },
    }
    result = fireworks.map_openai_params(
        {"response_format": response_format}, {}, "some_model", drop_params=False
    )
    assert result == {"response_format": response_format}


class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest):
    def get_base_audio_transcription_call_args(self) -> dict:
        return {
            "model": "fireworks_ai/whisper-v3",
            "api_base": "https://audio-prod.api.fireworks.ai/v1",
        }

    def get_custom_llm_provider(self) -> litellm.LlmProviders:
        return litellm.LlmProviders.FIREWORKS_AI


@pytest.mark.parametrize(
    "disable_add_transform_inline_image_block",
    [True, False],
)
def test_document_inlining_example(disable_add_transform_inline_image_block):
    """
    Document inlining appends ``#transform=inline`` to image/PDF URLs in the
    outgoing request unless explicitly disabled. Assert the transform on the
    serialized payload rather than making a live Fireworks call — the live
    call only proved the model responded and broke whenever Fireworks rotated
    its serverless model catalog.
    """
    from unittest.mock import patch

    from litellm import completion
    from litellm.llms.custom_httpx.http_handler import HTTPHandler

    client = HTTPHandler()
    pdf_url = "https://storage.googleapis.com/fireworks-public/test/sample_resume.pdf"

    with patch.object(client, "post") as mock_post:
        try:
            completion(
                model="fireworks_ai/accounts/fireworks/models/deepseek-v3p1",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image_url",
                                "image_url": {"url": pdf_url},
                            },
                            {
                                "type": "text",
                                "text": "What are the candidate's BA and MBA GPAs?",
                            },
                        ],
                    }
                ],
                disable_add_transform_inline_image_block=disable_add_transform_inline_image_block,
                client=client,
            )
        except Exception as e:
            print(e)

        mock_post.assert_called_once()
        json_data = json.loads(mock_post.call_args.kwargs["data"])
        sent_url = json_data["messages"][0]["content"][0]["image_url"]["url"]
        if disable_add_transform_inline_image_block is True:
            assert sent_url == pdf_url
            assert "#transform=inline" not in sent_url
        else:
            assert sent_url == pdf_url + "#transform=inline"


@pytest.mark.parametrize(
    "content, model, expected_url",
    [
        (
            {"image_url": "http://example.com/image.png"},
            "gpt-4",
            "http://example.com/image.png#transform=inline",
        ),
        (
            {"image_url": {"url": "http://example.com/image.png"}},
            "gpt-4",
            {"url": "http://example.com/image.png#transform=inline"},
        ),
        (
            {"image_url": "http://example.com/image.png"},
            "vision-gpt",
            "http://example.com/image.png",
        ),
        # data: URLs must never have #transform=inline appended — doing so
        # corrupts the base64 payload (fixes #23583).
        # URI schemes are case-insensitive (RFC 3986) so check all variants.
        (
            {"image_url": "data:image/png;base64,iVBORw0KGgo="},
            "gpt-4",
            "data:image/png;base64,iVBORw0KGgo=",
        ),
        (
            {"image_url": {"url": "data:image/jpeg;base64,/9j/4AAQ=="}},
            "gpt-4",
            {"url": "data:image/jpeg;base64,/9j/4AAQ=="},
        ),
        (
            {"image_url": "Data:image/png;base64,iVBORw0KGgo="},
            "gpt-4",
            "Data:image/png;base64,iVBORw0KGgo=",
        ),
    ],
)
def test_transform_inline(content, model, expected_url):

    result = litellm.FireworksAIConfig()._add_transform_inline_image_block(
        content=content, model=model, disable_add_transform_inline_image_block=False
    )
    if isinstance(expected_url, str):
        assert result["image_url"] == expected_url
    else:
        assert result["image_url"]["url"] == expected_url["url"]


@pytest.mark.parametrize(
    "model, is_disabled, expected_url",
    [
        ("gpt-4", True, "http://example.com/image.png"),
        ("vision-gpt", False, "http://example.com/image.png"),
        ("gpt-4", False, "http://example.com/image.png#transform=inline"),
    ],
)
def test_global_disable_flag(model, is_disabled, expected_url):
    content = {"image_url": "http://example.com/image.png"}
    result = litellm.FireworksAIConfig()._add_transform_inline_image_block(
        content=content,
        model=model,
        disable_add_transform_inline_image_block=is_disabled,
    )
    assert result["image_url"] == expected_url
    litellm.disable_add_transform_inline_image_block = False  # Reset for other tests


def test_global_disable_flag_with_transform_messages_helper(monkeypatch):
    from openai import OpenAI
    from unittest.mock import patch
    from litellm import completion
    from litellm.llms.custom_httpx.http_handler import HTTPHandler

    client = HTTPHandler()

    monkeypatch.setattr(litellm, "disable_add_transform_inline_image_block", True)

    with patch.object(
        client,
        "post",
    ) as mock_post:
        try:
            completion(
                model="fireworks_ai/accounts/fireworks/models/deepseek-v3p1",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": "What's in this image?"},
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": "https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/c233c9ade2ccb5491072ae232c814942.png"
                                },
                            },
                        ],
                    }
                ],
                client=client,
            )
        except Exception as e:
            print(e)

        mock_post.assert_called_once()
        print(mock_post.call_args.kwargs)
        json_data = json.loads(mock_post.call_args.kwargs["data"])
        assert (
            "#transform=inline"
            not in json_data["messages"][0]["content"][1]["image_url"]["url"]
        )
