#!/usr/bin/env python3
"""
Test to verify the Google GenAI transformation logic for generateContent parameters
"""
import os
import sys

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

import pytest

from litellm.llms.gemini.google_genai.transformation import GoogleGenAIConfig
from litellm.responses.litellm_completion_transformation.transformation import (
    LiteLLMCompletionResponsesConfig,
)


def test_map_generate_content_optional_params_response_json_schema_camelcase():
    """Test that responseJsonSchema (camelCase) is passed through correctly"""
    config = GoogleGenAIConfig()

    generate_content_config_dict = {
        "responseJsonSchema": {
            "type": "object",
            "properties": {"recipe_name": {"type": "string"}},
        },
        "temperature": 1.0,
    }

    result = config.map_generate_content_optional_params(
        generate_content_config_dict=generate_content_config_dict,
        model="gemini/gemini-3-flash-preview",
    )

    # responseJsonSchema should be in the result (camelCase format for Google GenAI API)
    assert "responseJsonSchema" in result
    assert (
        result["responseJsonSchema"]
        == generate_content_config_dict["responseJsonSchema"]
    )
    assert "temperature" in result
    assert result["temperature"] == 1.0


def test_map_generate_content_optional_params_response_schema_snakecase():
    """Test that response_schema (snake_case) is converted to responseJsonSchema (camelCase)"""
    config = GoogleGenAIConfig()

    generate_content_config_dict = {
        "response_json_schema": {
            "type": "object",
            "properties": {"recipe_name": {"type": "string"}},
        },
        "temperature": 1.0,
    }

    result = config.map_generate_content_optional_params(
        generate_content_config_dict=generate_content_config_dict,
        model="gemini/gemini-3-flash-preview",
    )

    # response_schema should be converted to responseJsonSchema (camelCase)
    assert "responseJsonSchema" in result
    assert (
        result["responseJsonSchema"]
        == generate_content_config_dict["response_json_schema"]
    )
    assert "temperature" in result


def test_map_generate_content_optional_params_thinking_config_camelcase():
    """Test that thinkingConfig (camelCase) is passed through correctly"""
    config = GoogleGenAIConfig()

    generate_content_config_dict = {
        "thinkingConfig": {"thinkingLevel": "minimal", "includeThoughts": True},
        "temperature": 1.0,
    }

    result = config.map_generate_content_optional_params(
        generate_content_config_dict=generate_content_config_dict,
        model="gemini/gemini-3-flash-preview",
    )

    # thinkingConfig should be in the result (camelCase format for Google GenAI API)
    assert "thinkingConfig" in result
    assert result["thinkingConfig"]["thinkingLevel"] == "minimal"
    assert result["thinkingConfig"]["includeThoughts"] is True
    assert "temperature" in result


def test_map_generate_content_optional_params_thinking_config_snakecase():
    """Test that thinking_config (snake_case) is converted to thinkingConfig (camelCase)"""
    config = GoogleGenAIConfig()

    generate_content_config_dict = {
        "thinking_config": {"thinkingLevel": "medium", "includeThoughts": True},
        "temperature": 1.0,
    }

    result = config.map_generate_content_optional_params(
        generate_content_config_dict=generate_content_config_dict,
        model="gemini/gemini-3-flash-preview",
    )

    # thinking_config should be converted to thinkingConfig (camelCase)
    assert "thinkingConfig" in result
    assert result["thinkingConfig"]["thinkingLevel"] == "medium"
    assert result["thinkingConfig"]["includeThoughts"] is True
    assert "thinking_config" not in result  # Should not be in snake_case format
    assert "temperature" in result


def test_map_generate_content_optional_params_mixed_formats():
    """Test that both camelCase and snake_case parameters work together"""
    config = GoogleGenAIConfig()

    generate_content_config_dict = {
        "responseJsonSchema": {
            "type": "object",
            "properties": {"recipe_name": {"type": "string"}},
        },
        "thinking_config": {"thinkingLevel": "low", "includeThoughts": True},
        "temperature": 1.0,
        "max_output_tokens": 100,
    }

    result = config.map_generate_content_optional_params(
        generate_content_config_dict=generate_content_config_dict,
        model="gemini/gemini-3-flash-preview",
    )

    # All parameters should be converted to camelCase
    assert "responseJsonSchema" in result
    assert "thinkingConfig" in result
    assert result["thinkingConfig"]["thinkingLevel"] == "low"
    assert "temperature" in result
    assert "maxOutputTokens" in result  # This one stays as-is if it's in supported list


def test_map_generate_content_optional_params_response_mime_type():
    """Test that responseMimeType is handled correctly"""
    config = GoogleGenAIConfig()

    generate_content_config_dict = {
        "responseMimeType": "application/json",
        "responseJsonSchema": {
            "type": "object",
            "properties": {"recipe_name": {"type": "string"}},
        },
    }

    result = config.map_generate_content_optional_params(
        generate_content_config_dict=generate_content_config_dict,
        model="gemini/gemini-3-flash-preview",
    )

    # responseMimeType should be passed through (it's already camelCase)
    assert "responseMimeType" in result or "response_mime_type" in result
    assert "responseJsonSchema" in result


def test_responses_api_reasoning_dict_format():
    """Test that reasoning parameter with dict format is mapped to reasoning_effort"""
    from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams

    responses_api_request: ResponsesAPIOptionalRequestParams = {
        "reasoning": {"effort": "high"},
        "temperature": 1.0,
    }

    result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
        model="gemini/2.5-pro",
        input="Hello, what is the capital of France?",
        responses_api_request=responses_api_request,
    )

    # reasoning_effort should be extracted from reasoning dict
    assert "reasoning_effort" in result
    assert result["reasoning_effort"] == "high"


def test_responses_api_reasoning_string_format():
    """Test that reasoning parameter with string format is mapped to reasoning_effort"""
    from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams

    responses_api_request: ResponsesAPIOptionalRequestParams = {
        "reasoning": "medium",  # Could be a string directly
        "temperature": 1.0,
    }

    result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
        model="gemini/2.5-pro",
        input="Hello, what is the capital of France?",
        responses_api_request=responses_api_request,
    )

    # reasoning_effort should be extracted from reasoning string
    assert "reasoning_effort" in result
    assert result["reasoning_effort"] == "medium"


def test_responses_api_reasoning_low_effort():
    """Test that low reasoning effort is correctly mapped"""
    from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams

    responses_api_request: ResponsesAPIOptionalRequestParams = {
        "reasoning": {"effort": "low"},
    }

    result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
        model="gemini/2.5-pro",
        input="Test",
        responses_api_request=responses_api_request,
    )

    assert "reasoning_effort" in result
    assert result["reasoning_effort"] == "low"


def test_responses_api_no_reasoning():
    """Test that no reasoning_effort is included when reasoning is not provided"""
    from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams

    responses_api_request: ResponsesAPIOptionalRequestParams = {
        "temperature": 1.0,
    }

    result = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request(
        model="gemini/2.5-pro",
        input="Test",
        responses_api_request=responses_api_request,
    )

    # reasoning_effort should not be in result if not provided (filtered out as None)
    assert "reasoning_effort" not in result or result.get("reasoning_effort") is None


def test_transform_generate_content_request_with_system_instruction():
    """Test that systemInstruction parameter is properly included in the request"""
    config = GoogleGenAIConfig()

    system_instruction = {"parts": [{"text": "You are a helpful assistant"}]}

    contents = [{"role": "user", "parts": [{"text": "Hello"}]}]

    generate_content_config_dict = {"temperature": 1.0, "maxOutputTokens": 100}

    # Call transform_generate_content_request
    result = config.transform_generate_content_request(
        model="gemini-3-flash-preview",
        contents=contents,
        tools=None,
        generate_content_config_dict=generate_content_config_dict,
        system_instruction=system_instruction,
    )

    # Verify that systemInstruction is in the request
    assert "systemInstruction" in result, "systemInstruction should be in request body"
    assert (
        result["systemInstruction"] == system_instruction
    ), "systemInstruction should match input"
    assert result["model"] == "gemini-3-flash-preview"
    assert result["contents"] == contents


def test_transform_generate_content_request_without_system_instruction():
    """Test that request works correctly without systemInstruction"""
    config = GoogleGenAIConfig()

    contents = [{"role": "user", "parts": [{"text": "Hello"}]}]

    generate_content_config_dict = {"temperature": 1.0}

    # Call transform_generate_content_request without system_instruction
    result = config.transform_generate_content_request(
        model="gemini-3-flash-preview",
        contents=contents,
        tools=None,
        generate_content_config_dict=generate_content_config_dict,
        system_instruction=None,
    )

    # Verify that systemInstruction is NOT in the request when not provided
    assert (
        "systemInstruction" not in result
    ), "systemInstruction should not be in request when None"
    assert result["model"] == "gemini-3-flash-preview"
    assert result["contents"] == contents


def test_transform_generate_content_request_system_instruction_with_tools():
    """Test that systemInstruction works correctly alongside tools"""
    config = GoogleGenAIConfig()

    system_instruction = {
        "parts": [{"text": "You are a helpful assistant that uses tools"}]
    }

    contents = [{"role": "user", "parts": [{"text": "What's the weather?"}]}]

    tools = [
        {
            "functionDeclarations": [
                {
                    "name": "get_weather",
                    "description": "Get weather information",
                    "parameters": {
                        "type": "object",
                        "properties": {"location": {"type": "string"}},
                    },
                }
            ]
        }
    ]

    generate_content_config_dict = {"temperature": 0.7}

    # Call transform_generate_content_request with both system_instruction and tools
    result = config.transform_generate_content_request(
        model="gemini-3-flash-preview",
        contents=contents,
        tools=tools,
        generate_content_config_dict=generate_content_config_dict,
        system_instruction=system_instruction,
    )

    # Verify that both systemInstruction and tools are in the request
    assert "systemInstruction" in result, "systemInstruction should be in request body"
    assert result["systemInstruction"] == system_instruction
    assert "tools" in result, "tools should be in request body"
    assert result["tools"] == tools
    assert result["model"] == "gemini-3-flash-preview"


def test_transform_generate_content_request_normalizes_response_schema_2_5():
    """For Gemini 2.0+, ``response_schema`` with ``$defs``/``$ref`` should be
    promoted to ``responseJsonSchema`` (which Gemini 2.0+ accepts natively),
    not forwarded as ``responseSchema`` (which rejects ``$defs``)."""
    config = GoogleGenAIConfig()

    schema = {
        "$defs": {
            "Highlight": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "detail": {"type": "string"},
                },
                "required": ["title"],
            }
        },
        "type": "object",
        "properties": {
            "park_name": {"type": "string"},
            "highlights": {
                "type": "array",
                "items": {"$ref": "#/$defs/Highlight"},
            },
        },
        "required": ["park_name", "highlights"],
    }

    result = config.transform_generate_content_request(
        model="gemini-2.5-flash-lite",
        contents=[{"role": "user", "parts": [{"text": "hi"}]}],
        tools=None,
        generate_content_config_dict={
            "responseMimeType": "application/json",
            "responseSchema": schema,
        },
        system_instruction=None,
    )

    gen_config = result["generationConfig"]
    assert "responseSchema" not in gen_config
    assert "responseJsonSchema" in gen_config
    normalized = gen_config["responseJsonSchema"]
    assert "$defs" in normalized
    assert normalized["properties"]["highlights"]["items"] == {
        "$ref": "#/$defs/Highlight"
    }


def test_transform_generate_content_request_flattens_response_schema_1_5():
    """For Gemini 1.5, ``responseSchema`` is kept but flattened via
    ``_build_vertex_schema`` so ``$defs``/``$ref`` are unpacked."""
    config = GoogleGenAIConfig()

    schema = {
        "$defs": {
            "Highlight": {
                "type": "object",
                "properties": {"title": {"type": "string"}},
                "required": ["title"],
            }
        },
        "type": "object",
        "properties": {
            "highlights": {
                "type": "array",
                "items": {"$ref": "#/$defs/Highlight"},
            }
        },
        "required": ["highlights"],
    }

    result = config.transform_generate_content_request(
        model="gemini-1.5-pro",
        contents=[{"role": "user", "parts": [{"text": "hi"}]}],
        tools=None,
        generate_content_config_dict={"responseSchema": schema},
        system_instruction=None,
    )

    gen_config = result["generationConfig"]
    assert "responseSchema" in gen_config
    assert "responseJsonSchema" not in gen_config
    normalized = gen_config["responseSchema"]
    assert "$defs" not in normalized
    items = normalized["properties"]["highlights"]["items"]
    assert "$ref" not in items
    assert "title" in items["properties"]
    assert items["properties"]["title"]["type"].lower() == "string"


def test_transform_generate_content_request_passes_through_response_json_schema():
    """If the caller already used ``responseJsonSchema``, it should be
    preserved (Gemini 2.0+ accepts standard JSON Schema as-is)."""
    config = GoogleGenAIConfig()

    schema = {
        "type": "object",
        "properties": {"name": {"type": "string"}},
        "required": ["name"],
    }

    result = config.transform_generate_content_request(
        model="gemini-2.5-flash-lite",
        contents=[{"role": "user", "parts": [{"text": "hi"}]}],
        tools=None,
        generate_content_config_dict={"responseJsonSchema": schema},
        system_instruction=None,
    )

    gen_config = result["generationConfig"]
    assert gen_config["responseJsonSchema"] == schema
    assert "responseSchema" not in gen_config


def test_transform_generate_content_request_preserves_response_json_schema_when_response_schema_co_present():
    """When both ``responseJsonSchema`` and ``responseSchema`` are supplied on
    Gemini 2.0+, the caller's ``responseJsonSchema`` must win — the
    ``responseSchema`` value must not clobber it."""
    config = GoogleGenAIConfig()

    caller_json_schema = {
        "type": "object",
        "properties": {"answer": {"type": "string"}},
        "required": ["answer"],
    }
    redundant_response_schema = {
        "type": "object",
        "properties": {"other": {"type": "string"}},
    }

    result = config.transform_generate_content_request(
        model="gemini-2.5-flash-lite",
        contents=[{"role": "user", "parts": [{"text": "hi"}]}],
        tools=None,
        generate_content_config_dict={
            "responseJsonSchema": caller_json_schema,
            "responseSchema": redundant_response_schema,
        },
        system_instruction=None,
    )

    gen_config = result["generationConfig"]
    assert "responseSchema" not in gen_config
    assert gen_config["responseJsonSchema"] == caller_json_schema


def test_response_schema_normalization_parity_across_chat_and_native_paths():
    """Parity guard between the two Gemini schema-normalization paths.

    The native ``generateContent`` path (``_normalize_response_schema``) must
    produce the same normalized schema as the ``/chat/completions`` path
    (``apply_response_schema_transformation``) for the same input schema, on
    both Gemini 2.0+ and Gemini 1.5. If either implementation drifts, this
    test fails — forcing both paths to be updated together.
    """
    from copy import deepcopy

    from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
        VertexGeminiConfig,
    )

    schema = {
        "$defs": {
            "Highlight": {
                "type": "object",
                "properties": {
                    "title": {"type": "string"},
                    "detail": {"type": "string"},
                },
                "required": ["title"],
            }
        },
        "type": "object",
        "properties": {
            "park_name": {"type": "string"},
            "highlights": {
                "type": "array",
                "items": {"$ref": "#/$defs/Highlight"},
            },
            "rating": {"anyOf": [{"type": "number"}, {"type": "null"}]},
        },
        "required": ["park_name", "highlights"],
    }

    native_config = GoogleGenAIConfig()
    chat_config = VertexGeminiConfig()

    for model, native_out_key, chat_out_key in (
        ("gemini-2.5-flash-lite", "responseJsonSchema", "response_json_schema"),
        ("gemini-1.5-pro", "responseSchema", "response_schema"),
    ):
        native_dict = {"responseSchema": deepcopy(schema)}
        native_config._normalize_response_schema(native_dict, model)

        chat_optional_params: dict = {}
        chat_config.apply_response_schema_transformation(
            value={"type": "json_schema", "response_schema": deepcopy(schema)},
            optional_params=chat_optional_params,
            model=model,
        )

        assert native_dict[native_out_key] == chat_optional_params[chat_out_key], (
            f"Schema normalization drifted between native generateContent and "
            f"chat/completions paths for {model}. Update both "
            f"GoogleGenAIConfig._normalize_response_schema and "
            f"VertexGeminiConfig.apply_response_schema_transformation together."
        )


def test_transform_generate_content_request_without_schema_unchanged():
    """No schema in config → no normalization side effects."""
    config = GoogleGenAIConfig()

    result = config.transform_generate_content_request(
        model="gemini-2.5-flash-lite",
        contents=[{"role": "user", "parts": [{"text": "hi"}]}],
        tools=None,
        generate_content_config_dict={"temperature": 0.7},
        system_instruction=None,
    )

    assert result["generationConfig"]["temperature"] == 0.7
    assert "responseSchema" not in result["generationConfig"]
    assert "responseJsonSchema" not in result["generationConfig"]


def test_validate_environment_with_dict_api_key():
    """
    Test that validate_environment correctly handles api_key as a dict.

    This happens when using custom api_base with Gemini - the auth_header
    is returned as {"x-goog-api-key": "sk-test"} and should be merged into
    headers instead of being set as a header value.

    Regression test for: https://github.com/BerriAI/litellm/issues/xxxxx
    """
    config = GoogleGenAIConfig()

    # Simulate the case where auth_header is a dict (custom api_base scenario)
    auth_header_dict = {"x-goog-api-key": "sk-test-key-123"}

    result = config.validate_environment(
        api_key=auth_header_dict,
        headers=None,
        model="gemini-2.5-pro",
        litellm_params={},
    )

    # The dict should be merged into headers, not set as a value
    assert "x-goog-api-key" in result, "x-goog-api-key should be in headers"
    assert (
        result["x-goog-api-key"] == "sk-test-key-123"
    ), "API key should be the string value, not a dict"
    assert isinstance(
        result["x-goog-api-key"], str
    ), "Header value should be a string, not a dict"
    assert "Content-Type" in result, "Content-Type should be in headers"
    assert result["Content-Type"] == "application/json"


def test_validate_environment_with_string_api_key():
    """
    Test that validate_environment correctly handles api_key as a string.

    This is the normal case when using standard Gemini API.
    """
    config = GoogleGenAIConfig()

    # Normal case: api_key is a string
    api_key_string = "sk-test-key-456"

    result = config.validate_environment(
        api_key=api_key_string, headers=None, model="gemini-2.5-pro", litellm_params={}
    )

    # The string should be set as the header value
    assert "x-goog-api-key" in result, "x-goog-api-key should be in headers"
    assert result["x-goog-api-key"] == "sk-test-key-456", "API key should match input"
    assert isinstance(result["x-goog-api-key"], str), "Header value should be a string"
    assert "Content-Type" in result, "Content-Type should be in headers"


def test_validate_environment_with_extra_headers():
    """
    Test that validate_environment correctly merges extra headers with dict api_key.
    """
    config = GoogleGenAIConfig()

    # Custom api_base scenario with additional headers
    auth_header_dict = {"x-goog-api-key": "sk-test-key-789"}
    extra_headers = {"X-Custom-Header": "custom-value"}

    result = config.validate_environment(
        api_key=auth_header_dict,
        headers=extra_headers,
        model="gemini-2.5-pro",
        litellm_params={},
    )

    # Both the auth dict and extra headers should be merged
    assert "x-goog-api-key" in result, "x-goog-api-key should be in headers"
    assert (
        result["x-goog-api-key"] == "sk-test-key-789"
    ), "API key should be correctly set"
    assert isinstance(result["x-goog-api-key"], str), "Header value should be a string"
    assert "X-Custom-Header" in result, "Extra headers should be merged"
    assert result["X-Custom-Header"] == "custom-value"
    assert "Content-Type" in result
