import json

from aistudio_api.infrastructure.gateway.wire_codec import (
    AistudioWireCodec,
    build_image_generation_search_tool,
    build_tools_from_names,
    modify_body,
)


def test_modify_body_updates_generation_config_and_prompt():
    original = '["models/original",[[[[null,"old"]],"user"]],null,[null,null,null,128,0.5,0.8,16],"!snap",null,null]'
    rewritten = modify_body(
        original,
        model="models/new",
        prompt="new prompt",
        system_instruction="sys",
        max_tokens=256,
        temperature=0.2,
        top_p=0.9,
        top_k=32,
    )

    assert '"models/new"' in rewritten
    assert '"new prompt"' in rewritten
    assert '"sys"' in rewritten
    body = json.loads(rewritten)
    assert body[3][3] == 256
    assert body[3][4] == 0.2
    assert body[3][5] == 0.9
    assert body[3][6] == 32
    assert body[3][16] == [1, None, None, 3]
    assert len(body[3]) <= 17 or body[3][17] is None


def test_build_image_generation_search_tool_variants():
    assert build_image_generation_search_tool(google_search=True, image_search=False) == [None, None, None, [None, [[]]]]
    assert build_image_generation_search_tool(google_search=False, image_search=True) == [None, None, None, [None, [None, []]]]
    assert build_image_generation_search_tool(google_search=True, image_search=True) == [None, None, None, [None, [[], []]]]
    assert build_image_generation_search_tool(google_search=False, image_search=False) is None


def test_build_tools_from_names_supports_image_tool_aliases():
    assert build_tools_from_names(["google_search"], model="models/gemini-3.1-flash-image-preview", is_image_model=True) == [[None, None, None, [None, [[]]]]]
    assert build_tools_from_names(["image_search"], model="models/gemini-3.1-flash-image-preview", is_image_model=True) == [[None, None, None, [None, [None, []]]]]
    assert build_tools_from_names(["google_search_and_image_search"], model="models/gemini-3.1-flash-image-preview", is_image_model=True) == [
        [None, None, None, [None, [[], []]]]
    ]


def test_build_tools_from_names_merges_image_search_flags_into_single_tool():
    assert build_tools_from_names(
        ["google_search", "image_search"],
        model="models/gemini-3.1-flash-image-preview",
        is_image_model=True,
    ) == [[None, None, None, [None, [[], []]]]]


def test_build_tools_from_names_restricts_gemma_builtin_tools():
    assert build_tools_from_names(["google_search", "code_execution"], model="models/gemma-4-31b-it") == [
        [None, None, None, [None, [[]]]],
        [[]],
    ]

    try:
        build_tools_from_names(["google_maps"], model="models/gemma-4-31b-it")
    except ValueError as exc:
        assert "not allowed" in str(exc)
    else:
        raise AssertionError("expected ValueError")


def test_build_tools_from_names_allows_four_gemini_builtin_tools():
    assert build_tools_from_names(
        ["google_search", "code_execution", "google_maps", "url_context"],
        model="models/gemini-3.5-flash",
    ) == [
        [None, None, None, [None, [[]]]],
        [[]],
        [None, None, None, None, None, None, None, None, None, None, []],
        [None, None, None, None, None, None, None, []],
    ]


def test_build_tools_from_names_restricts_unknown_model_to_safe_subset():
    try:
        build_tools_from_names(["google_maps"], model="models/learnlm-1.5-pro-experimental")
    except ValueError as exc:
        assert "not allowed" in str(exc)
    else:
        raise AssertionError("expected ValueError")


def test_wire_codec_decodes_semantic_fields():
    codec = AistudioWireCodec()
    original = '["models/original",[[[[null,"old"]],"user"]],null,[null,null,null,128,0.5,0.8,16],"!snap",[[[null,"sys"]],"user"],[[[]]]]'

    decoded = codec.decode(original)

    assert decoded.model == "models/original"
    assert decoded.snapshot == "!snap"
    assert decoded.contents[0].role == "user"
    assert decoded.contents[0].parts[-1].text == "old"
    assert decoded.system_instruction is not None
    assert decoded.system_instruction.parts[0].text == "sys"
    assert decoded.generation_config.max_tokens == 128


def test_modify_body_sanitizes_plain_text_generation_config():
    original = '["models/original",[[[[null,"old"]],"user"]],null,[null,null,null,128,0.5,0.8,16,"application/json",[6],null,null,null,null,null,null,null,[1,null,null,3]],"!snap",null,null]'
    rewritten = modify_body(
        original,
        model="models/gemma-4-31b-it",
        prompt="hello",
    )

    assert '"text/plain"' in rewritten
    assert '"application/json"' not in rewritten
    assert '[6]' not in rewritten
    assert json.loads(rewritten)[3][16] == [1, None, None, 3]
    assert len(json.loads(rewritten)[3]) <= 17 or json.loads(rewritten)[3][17] is None


def test_modify_body_enables_thinking_for_any_model():
    original = '["models/original",[[[[null,"old"]],"user"]],null,[null,null,null,128,0.5,0.8,16],"!snap",null,null]'
    rewritten = modify_body(
        original,
        model="models/new-text-model",
        prompt="hello",
    )

    body = json.loads(rewritten)
    assert body[3][16] == [1, None, None, 3]
    assert len(body[3]) <= 17 or body[3][17] is None


def test_modify_body_can_override_image_output_resolution():
    original = json.dumps(
        [
            "models/original",
            [[[[None, "old"]], "user"]],
            None,
            [None] * 27,
            "!snap",
            None,
            None,
        ]
    )
    rewritten = modify_body(
        original,
        model="models/gemini-3.1-flash-image-preview",
        prompt="hello",
        generation_config_overrides={"output_resolution": ["16:9", "1K"]},
    )

    body = json.loads(rewritten)
    assert body[3][26] == ["16:9", "1K"]


def test_modify_body_applies_model_safety_defaults():
    original = '["models/original",[[[[null,"old"]],"user"]],null,[null,null,null,128,0.5,0.8,16],"!snap",null,null]'
    rewritten = modify_body(
        original,
        model="models/gemma-4-31b-it",
        prompt="hello",
    )

    body = json.loads(rewritten)
    assert body[2] == [
        [None, None, 7, 5],
        [None, None, 8, 5],
        [None, None, 9, 5],
        [None, None, 10, 5],
    ]


def test_modify_body_applies_image_model_defaults():
    original = json.dumps(
        [
            "models/original",
            [[[[None, "old"]], "user"]],
            [[None, None, 7, 4]],
            [None] * 27,
            "!snap",
            None,
            None,
        ]
    )
    rewritten = modify_body(
        original,
        model="models/gemini-3.1-flash-image-preview",
        prompt="hello",
    )

    body = json.loads(rewritten)
    assert body[2] is None
    assert body[3][14] == [2]
    assert body[3][16] == [1, None, None, 4]


def test_modify_body_keeps_image_model_tools():
    original = json.dumps(
        [
            "models/original",
            [[[[None, "old"]], "user"]],
            None,
            [None] * 27,
            "!snap",
            None,
            None,
        ]
    )
    rewritten = modify_body(
        original,
        model="models/gemini-3.1-flash-image-preview",
        prompt="hello",
        tools=[[None, None, None, [None, [[], []]]]],
    )

    body = json.loads(rewritten)
    assert body[6] == [[None, None, None, [None, [[], []]]]]


def test_modify_body_safety_off_uses_off_threshold():
    original = '["models/original",[[[[null,"old"]],"user"]],null,[null,null,null,128,0.5,0.8,16],"!snap",null,null]'
    rewritten = modify_body(
        original,
        model="models/gemma-4-31b-it",
        prompt="hello",
        safety_off=True,
    )

    body = json.loads(rewritten)
    assert body[2] == [
        [None, None, 7, 5],
        [None, None, 8, 5],
        [None, None, 9, 5],
        [None, None, 10, 5],
    ]


def test_modify_body_can_override_safety_settings():
    original = '["models/original",[[[[null,"old"]],"user"]],null,[null,null,null,128,0.5,0.8,16],"!snap",null,null]'
    rewritten = modify_body(
        original,
        model="models/gemini-3.5-flash",
        prompt="hello",
        safety_settings=[
            [None, None, 7, 1],
            [None, None, 8, 4],
            [None, None, 9, 3],
            [None, None, 10, 2],
        ],
    )

    body = json.loads(rewritten)
    assert body[2] == [
        [None, None, 7, 1],
        [None, None, 8, 4],
        [None, None, 9, 3],
        [None, None, 10, 2],
    ]


def test_modify_body_keeps_structured_generation_config_for_gemini_mode():
    original = '["models/original",[[[[null,"old"]],"user"]],null,[null,["6"],null,128,0.5,0.8,16,"application/json",[6],0.1,0.2,true,5,null,[2,1],null,[1,null,null,3],2],"!snap",null,null,null,null,null,1,"cached",null,[[null,null,"Asia/Shanghai"],null,1]]'
    rewritten = modify_body(
        original,
        model="models/gemini-2.5-pro-preview-05-06",
        prompt="hello",
        generation_config_overrides={
            "stop_sequences": ["STOP"],
            "response_mime_type": "application/json",
            "response_schema": [6],
            "presence_penalty": 0.3,
            "frequency_penalty": 0.4,
            "response_logprobs": True,
            "logprobs": 7,
            "image_output_mode": [2, 1],
            "thinking_config": [1, None, None, 3],
            "media_resolution": 2,
        },
        sanitize_plain_text=False,
    )

    body = json.loads(rewritten)
    assert body[3][1] == ["STOP"]
    assert body[3][7] == "application/json"
    assert body[3][8] == [6]
    assert body[3][9] == 0.3
    assert body[3][10] == 0.4
    assert body[3][11] is True
    assert body[3][12] == 7
    assert body[3][14] == [2, 1]
    assert body[3][16] == [1, None, None, 3]
    assert body[3][17] == 2