import pytest

from aistudio_api.api.schemas import GeminiContent, GeminiGenerateContentRequest, GeminiGenerationConfig, GeminiPart
from aistudio_api.application.chat_service import normalize_gemini_request
from aistudio_api.infrastructure.gateway.wire_types import AistudioImageOutputMode


def test_normalize_gemini_request_exposes_generation_config_overrides():
    req = GeminiGenerateContentRequest(
        contents=[GeminiContent(role="user", parts=[GeminiPart(text="hello")])],
        generationConfig=GeminiGenerationConfig(
            stopSequences=["6"],
            temperature=1,
            topP=0.95,
            topK=64,
            maxOutputTokens=65536,
            responseMimeType="text/plain",
            responseSchema={
                "type": "object",
                "properties": {"test_response": {"type": "string"}},
                "propertyOrdering": ["test_response"],
            },
            presencePenalty=0.1,
            frequencyPenalty=0.2,
            responseLogprobs=True,
            logprobs=5,
            mediaResolution=2,
            thinkingConfig=[1, None, None, 3],
        ),
    )

    normalized = normalize_gemini_request(req, "models/gemini-3.1-flash-image-preview")

    assert normalized["generation_config_overrides"] == {
        "stop_sequences": ["6"],
        "max_tokens": 65536,
        "temperature": 1,
        "top_p": 0.95,
        "top_k": 64,
        "response_mime_type": "text/plain",
        "response_schema": [6, None, None, None, None, None, [["test_response", [1]]], None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, ["test_response"]],
        "presence_penalty": 0.1,
        "frequency_penalty": 0.2,
        "response_logprobs": True,
        "logprobs": 5,
        "image_output_mode": AistudioImageOutputMode.image_only(),
        "media_resolution": 2,
        "thinking_config": [1, None, None, 3],
    }


def test_normalize_gemini_request_maps_official_image_generation_fields():
    req = GeminiGenerateContentRequest.model_validate(
        {
            "contents": [
                {
                    "role": "user",
                    "parts": [
                        {"inlineData": {"mimeType": "image/jpeg", "data": "/9j/4AAQSkZJRgABAQAA...."}},
                        {"text": "INSERT_INPUT_HERE"},
                    ],
                }
            ],
            "generationConfig": {
                "responseModalities": ["IMAGE", "TEXT"],
                "thinkingConfig": {"thinkingLevel": "HIGH"},
                "imageConfig": {
                    "aspectRatio": "9:16",
                    "imageSize": "4K",
                    "personGeneration": "",
                },
            },
            "tools": [
                {
                    "googleSearch": {
                        "searchTypes": {
                            "webSearch": {},
                            "imageSearch": {},
                        }
                    }
                }
            ],
        }
    )

    normalized = normalize_gemini_request(req, "models/gemini-3.1-flash-image-preview")

    assert normalized["tools"] == [[None, None, None, [None, [[], []]]]]
    assert normalized["generation_config_overrides"] == {
        "image_output_mode": AistudioImageOutputMode.text_and_image(),
        "thinking_config": [1, None, None, 3],
        "output_resolution": ["9:16", "4K"],
    }
    assert normalized["capture_images"] is not None
    assert len(normalized["capture_images"]) == 1
    assert normalized["contents"][0].parts[0].inline_data == ("image/jpeg", "/9j/4AAQSkZJRgABAQAA....")


def test_normalize_gemini_request_encodes_function_declarations_to_wire_tools():
    req = GeminiGenerateContentRequest(
        contents=[GeminiContent(role="user", parts=[GeminiPart(text="hello")])],
        tools=[
            {
                "functionDeclarations": [
                    {
                        "name": "getWeather",
                        "description": "gets the weather for a requested city",
                        "parameters": {
                            "type": "object",
                            "properties": {"city": {"type": "string"}},
                            "propertyOrdering": ["city"],
                        },
                    }
                ]
            }
        ],
    )

    normalized = normalize_gemini_request(req, "models/gemma-4-31b-it")

    assert normalized["tools"] == [
        [
            None,
            [
                [
                    "getWeather",
                    "gets the weather for a requested city",
                    [6, None, None, None, None, None, [["city", [1]]], None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, ["city"]],
                ]
            ],
        ],
        [None, None, None, [None, [[]]]],
    ]


def test_normalize_gemini_request_applies_gemma_default_tools():
    req = GeminiGenerateContentRequest(
        contents=[GeminiContent(role="user", parts=[GeminiPart(text="hello")])],
    )

    normalized = normalize_gemini_request(req, "models/gemma-4-31b-it")

    assert normalized["tools"] == [[None, None, None, [None, [[]]]]]


def test_normalize_gemini_request_encodes_builtin_tools_to_wire():
    req = GeminiGenerateContentRequest(
        contents=[GeminiContent(role="user", parts=[GeminiPart(text="hello")])],
        tools=[
            {
                "googleSearch": {},
                "googleMaps": {},
                "urlContext": {},
                "codeExecution": {},
            }
        ],
    )

    normalized = normalize_gemini_request(req, "models/gemini-3.5-flash")

    assert normalized["tools"] == [
        [[]],
        [None, None, None, [None, [[]]]],
        [None, None, None, None, None, None, None, None, None, None, []],
        [None, None, None, None, None, None, None, []],
    ]


def test_normalize_gemini_request_rejects_gemma_unsupported_builtin_tool():
    req = GeminiGenerateContentRequest(
        contents=[GeminiContent(role="user", parts=[GeminiPart(text="hello")])],
        tools=[
            {
                "googleMaps": {},
            }
        ],
    )

    with pytest.raises(ValueError, match="not allowed"):
        normalize_gemini_request(req, "models/gemma-4-31b-it")


def test_normalize_gemini_request_empty_tools_disables_model_defaults():
    req = GeminiGenerateContentRequest(
        contents=[GeminiContent(role="user", parts=[GeminiPart(text="hello")])],
        tools=[],
    )

    normalized = normalize_gemini_request(req, "models/gemma-4-31b-it")

    assert normalized["tools"] == []


def test_normalize_gemini_request_rejects_unsupported_person_generation():
    req = GeminiGenerateContentRequest.model_validate(
        {
            "contents": [{"role": "user", "parts": [{"text": "hello"}]}],
            "generationConfig": {
                "imageConfig": {
                    "personGeneration": "ALLOW_ADULT",
                }
            },
        }
    )

    with pytest.raises(ValueError, match="personGeneration"):
        normalize_gemini_request(req, "models/gemini-3.1-flash-image-preview")


def test_normalize_gemini_request_maps_official_text_model_fields():
    req = GeminiGenerateContentRequest.model_validate(
        {
            "contents": [
                {
                    "role": "user",
                    "parts": [
                        {"text": "INSERT_INPUT_HERE"},
                    ],
                }
            ],
            "generationConfig": {
                "thinkingConfig": {"thinkingLevel": "HIGH"},
            },
            "safetySettings": [
                {
                    "category": "HARM_CATEGORY_HARASSMENT",
                    "threshold": "BLOCK_LOW_AND_ABOVE",
                },
                {
                    "category": "HARM_CATEGORY_HATE_SPEECH",
                    "threshold": "BLOCK_NONE",
                },
                {
                    "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
                    "threshold": "BLOCK_ONLY_HIGH",
                },
                {
                    "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
                    "threshold": "BLOCK_MEDIUM_AND_ABOVE",
                },
            ],
            "tools": [
                {"urlContext": {}},
                {"codeExecution": {}},
                {"googleSearch": {}},
            ],
        }
    )

    normalized = normalize_gemini_request(req, "models/gemini-3.5-flash")

    assert normalized["tools"] == [
        [None, None, None, None, None, None, None, []],
        [[]],
        [None, None, None, [None, [[]]]],
    ]
    assert normalized["generation_config_overrides"] == {
        "thinking_config": [1, None, None, 3],
    }
    assert normalized["safety_settings"] == [
        [None, None, 7, 1],
        [None, None, 8, 4],
        [None, None, 9, 3],
        [None, None, 10, 2],
    ]