import asyncio
from unittest.mock import AsyncMock, patch

import pytest

from src import nodes


class FakeImageObject:
    def __init__(self, b64_json=None, url=None):
        self.b64_json = b64_json
        self.url = url


class FakeImageResponse:
    def __init__(self, data):
        self.data = data


def _text_source(model="imagen-4.0-generate-001", **overrides):
    settings = {"model": model, "size": "1024x1024", "num_images": 1}
    settings.update(overrides)
    return {"source": "text", "provider": {"provider": "gemini", "model": settings}}


def _image_source(model="gemini-2.5-flash-image", **overrides):
    settings = {"model": model, "size": "1024x1024", "num_images": 1}
    settings.update(overrides)
    return {"source": "image", "provider": {"provider": "gemini", "model": settings}}


def test_schema_lists_only_nano_banana_for_image_source():
    schema = nodes.APIImageGenerate.define_schema()
    source_input = next(i for i in schema.inputs if i.name == "source")
    text_option, image_option = source_input.options

    text_provider_input = text_option.inputs[0]
    gemini_text_option = next(o for o in text_provider_input.options if o.key == "gemini")
    text_model_ids = {o.key for o in gemini_text_option.inputs[0].options}
    assert "imagen-4.0-generate-001" in text_model_ids
    assert "gemini-2.5-flash-image" in text_model_ids

    image_provider_input = image_option.inputs[0]
    gemini_image_option = next(o for o in image_provider_input.options if o.key == "gemini")
    image_model_ids = {o.key for o in gemini_image_option.inputs[0].options}
    assert "gemini-2.5-flash-image" in image_model_ids
    assert "imagen-4.0-generate-001" not in image_model_ids

    for model_option in gemini_image_option.inputs[0].options:
        field_names = {i.name for i in model_option.inputs}
        assert {"size", "num_images", "reference_image"} <= field_names


def test_text_to_image_calls_aimage_generation():
    async def run():
        gen_mock = AsyncMock(
            return_value=FakeImageResponse(data=[FakeImageObject(b64_json="ZmFrZS1wbmc=")])
        )

        with patch.object(nodes, "aimage_generation", gen_mock), patch.object(
            nodes.Model,
            "image_bytes_list_to_tensor",
            return_value="fake-tensor",
        ) as tensor_mock:
            result = await nodes.APIImageGenerate.execute(
                prompt="a cat wearing a hat",
                source=_text_source(),
            )

        gen_mock.assert_awaited_once()
        assert gen_mock.call_args.kwargs["model"] == "gemini/imagen-4.0-generate-001"
        assert gen_mock.call_args.kwargs["prompt"] == "a cat wearing a hat"
        assert gen_mock.call_args.kwargs["size"] == "1024x1024"
        assert gen_mock.call_args.kwargs["n"] == 1
        tensor_mock.assert_called_once_with([b"fake-png"])
        assert result.args[0] == "fake-tensor"

    asyncio.run(run())


def test_image_to_image_calls_aimage_edit_with_batched_images():
    async def run():
        edit_mock = AsyncMock(
            return_value=FakeImageResponse(data=[FakeImageObject(b64_json="ZmFrZS1wbmc=")])
        )
        fake_tensor = object()
        fake_buffers = [object(), object()]

        source = _image_source()
        source["provider"]["model"]["reference_image"] = fake_tensor

        with patch.object(nodes, "aimage_edit", edit_mock), patch.object(
            nodes.GeminiModel,
            "image_tensor_to_bytesio_list",
            return_value=fake_buffers,
        ) as split_mock, patch.object(
            nodes.Model, "image_bytes_list_to_tensor", return_value="fake-tensor"
        ):
            result = await nodes.APIImageGenerate.execute(
                prompt="combine these two",
                source=source,
            )

        split_mock.assert_called_once_with(fake_tensor)
        edit_mock.assert_awaited_once()
        assert edit_mock.call_args.kwargs["model"] == "gemini/gemini-2.5-flash-image"
        assert edit_mock.call_args.kwargs["image"] == fake_buffers
        assert edit_mock.call_args.kwargs["prompt"] == "combine these two"
        assert result.args[0] == "fake-tensor"

    asyncio.run(run())


def test_image_to_image_requires_reference_image():
    async def run():
        edit_mock = AsyncMock()

        with patch.object(nodes, "aimage_edit", edit_mock):
            with pytest.raises(ValueError, match="reference_image"):
                await nodes.APIImageGenerate.execute(
                    prompt="anything",
                    source=_image_source(),
                )

        edit_mock.assert_not_awaited()

    asyncio.run(run())