"""Integration tests for bulkgen.providers (image and text).

Mock-heavy tests produce many Any-typed expressions from MagicMock.
"""
# pyright: reportAny=false, reportUnknownMemberType=false

from __future__ import annotations

import base64
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from bulkgen.config import TargetConfig
from bulkgen.providers.bfl import BFLResult
from bulkgen.providers.blackforest import BlackForestProvider
from bulkgen.providers.blackforest import (
    _encode_image_b64 as encode_image_b64,  # pyright: ignore[reportPrivateUsage]
)
from bulkgen.providers.mistral import MistralProvider
from bulkgen.providers.models import ModelInfo
from bulkgen.providers.openai_image import OpenAIImageProvider
from bulkgen.providers.registry import get_all_models


def _model(name: str) -> ModelInfo:
    """Look up a ModelInfo by name."""
    for m in get_all_models():
        if m.name == name:
            return m
    msg = f"Unknown test model: {name}"
    raise ValueError(msg)


def _make_bfl_mocks(
    image_bytes: bytes,
) -> tuple[BFLResult, MagicMock]:
    """Return (bfl_result, mock_http) for BFL image generation tests."""
    bfl_result = BFLResult(
        task_id="test-task-id", sample_url="https://example.com/img.png"
    )

    mock_response = MagicMock()
    mock_response.content = image_bytes
    mock_response.raise_for_status.return_value = None

    mock_http = AsyncMock()
    mock_http.get.return_value = mock_response
    mock_http.__aenter__ = AsyncMock(return_value=mock_http)
    mock_http.__aexit__ = AsyncMock(return_value=False)

    return bfl_result, mock_http


def _make_mistral_mock(response: MagicMock) -> AsyncMock:
    """Return a mock Mistral client."""
    mock_client = AsyncMock()
    mock_client.chat.complete_async.return_value = response
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=False)
    return mock_client


def _make_text_response(content: str | None) -> MagicMock:
    """Return a mock Mistral response with one choice."""
    choice = MagicMock()
    choice.message.content = content
    response = MagicMock()
    response.choices = [choice]
    return response


class TestBlackForestProvider:
    """Test BlackForestProvider with mocked BFL client and HTTP."""

    @pytest.fixture
    def image_bytes(self) -> bytes:
        return b"\x89PNG\r\n\x1a\n" + b"\x00" * 100

    async def test_basic_image_generation(
        self, project_dir: Path, image_bytes: bytes
    ) -> None:
        target_config = TargetConfig(prompt="A red square")
        bfl_result, mock_http = _make_bfl_mocks(image_bytes)

        with (
            patch("bulkgen.providers.blackforest.BFLClient") as mock_cls,
            patch("bulkgen.providers.blackforest.httpx.AsyncClient") as mock_http_cls,
        ):
            mock_cls.return_value.generate = AsyncMock(return_value=bfl_result)
            mock_http_cls.return_value = mock_http

            provider = BlackForestProvider(api_key="test-key")
            await provider.generate(
                target_name="out.png",
                target_config=target_config,
                resolved_prompt="A red square",
                resolved_model=_model("flux-pro-1.1"),
                project_dir=project_dir,
            )

        output = project_dir / "out.png"
        assert output.exists()
        assert output.read_bytes() == image_bytes

    async def test_image_with_dimensions(
        self, project_dir: Path, image_bytes: bytes
    ) -> None:
        target_config = TargetConfig(prompt="A banner", width=1920, height=480)
        bfl_result, mock_http = _make_bfl_mocks(image_bytes)

        with (
            patch("bulkgen.providers.blackforest.BFLClient") as mock_cls,
            patch("bulkgen.providers.blackforest.httpx.AsyncClient") as mock_http_cls,
        ):
            mock_generate = AsyncMock(return_value=bfl_result)
            mock_cls.return_value.generate = mock_generate
            mock_http_cls.return_value = mock_http

            provider = BlackForestProvider(api_key="test-key")
            await provider.generate(
                target_name="banner.png",
                target_config=target_config,
                resolved_prompt="A banner",
                resolved_model=_model("flux-pro-1.1"),
                project_dir=project_dir,
            )

            call_args = mock_generate.call_args
            inputs = call_args[0][1]
            assert inputs["width"] == 1920
            assert inputs["height"] == 480

    async def test_image_with_reference_image_flux1(
        self, project_dir: Path, image_bytes: bytes
    ) -> None:
        ref_path = project_dir / "ref.png"
        _ = ref_path.write_bytes(b"reference image data")

        target_config = TargetConfig(prompt="Like this", reference_images=["ref.png"])
        bfl_result, mock_http = _make_bfl_mocks(image_bytes)

        with (
            patch("bulkgen.providers.blackforest.BFLClient") as mock_cls,
            patch("bulkgen.providers.blackforest.httpx.AsyncClient") as mock_http_cls,
        ):
            mock_generate = AsyncMock(return_value=bfl_result)
            mock_cls.return_value.generate = mock_generate
            mock_http_cls.return_value = mock_http

            provider = BlackForestProvider(api_key="test-key")
            await provider.generate(
                target_name="out.png",
                target_config=target_config,
                resolved_prompt="Like this",
                resolved_model=_model("flux-pro-1.1"),
                project_dir=project_dir,
            )

            call_args = mock_generate.call_args
            inputs = call_args[0][1]
            assert "image_prompt" in inputs
            assert inputs["image_prompt"] == encode_image_b64(ref_path)

    async def test_image_with_reference_images_flux2(
        self, project_dir: Path, image_bytes: bytes
    ) -> None:
        ref1 = project_dir / "ref1.png"
        ref2 = project_dir / "ref2.png"
        _ = ref1.write_bytes(b"ref1 data")
        _ = ref2.write_bytes(b"ref2 data")

        target_config = TargetConfig(
            prompt="Combine", reference_images=["ref1.png", "ref2.png"]
        )
        bfl_result, mock_http = _make_bfl_mocks(image_bytes)

        with (
            patch("bulkgen.providers.blackforest.BFLClient") as mock_cls,
            patch("bulkgen.providers.blackforest.httpx.AsyncClient") as mock_http_cls,
        ):
            mock_generate = AsyncMock(return_value=bfl_result)
            mock_cls.return_value.generate = mock_generate
            mock_http_cls.return_value = mock_http

            provider = BlackForestProvider(api_key="test-key")
            await provider.generate(
                target_name="out.png",
                target_config=target_config,
                resolved_prompt="Combine",
                resolved_model=_model("flux-2-pro"),
                project_dir=project_dir,
            )

            call_args = mock_generate.call_args
            inputs = call_args[0][1]
            assert inputs["input_image"] == encode_image_b64(ref1)
            assert inputs["input_image_2"] == encode_image_b64(ref2)

    async def test_image_with_reference_image_kontext(
        self, project_dir: Path, image_bytes: bytes
    ) -> None:
        ref_path = project_dir / "ref.png"
        _ = ref_path.write_bytes(b"reference image data")

        target_config = TargetConfig(prompt="Edit", reference_images=["ref.png"])
        bfl_result, mock_http = _make_bfl_mocks(image_bytes)

        with (
            patch("bulkgen.providers.blackforest.BFLClient") as mock_cls,
            patch("bulkgen.providers.blackforest.httpx.AsyncClient") as mock_http_cls,
        ):
            mock_generate = AsyncMock(return_value=bfl_result)
            mock_cls.return_value.generate = mock_generate
            mock_http_cls.return_value = mock_http

            provider = BlackForestProvider(api_key="test-key")
            await provider.generate(
                target_name="out.png",
                target_config=target_config,
                resolved_prompt="Edit",
                resolved_model=_model("flux-kontext-pro"),
                project_dir=project_dir,
            )

            call_args = mock_generate.call_args
            inputs = call_args[0][1]
            assert inputs["input_image"] == encode_image_b64(ref_path)

    async def test_image_no_sample_url_raises(self, project_dir: Path) -> None:
        target_config = TargetConfig(prompt="x")

        with patch("bulkgen.providers.blackforest.BFLClient") as mock_cls:
            from bulkgen.providers.bfl import BFLError

            mock_cls.return_value.generate = AsyncMock(
                side_effect=BFLError("BFL task test ready but no sample URL: {}")
            )

            provider = BlackForestProvider(api_key="test-key")
            with pytest.raises(BFLError, match="no sample URL"):
                await provider.generate(
                    target_name="fail.png",
                    target_config=target_config,
                    resolved_prompt="x",
                    resolved_model=_model("flux-pro-1.1"),
                    project_dir=project_dir,
                )

    def test_encode_image_b64(self, project_dir: Path) -> None:
        data = b"test image bytes"
        f = project_dir / "test.png"
        _ = f.write_bytes(data)

        encoded = encode_image_b64(f)
        assert base64.b64decode(encoded) == data


class TestMistralProvider:
    """Test MistralProvider with mocked Mistral client."""

    async def test_basic_text_generation(self, project_dir: Path) -> None:
        target_config = TargetConfig(prompt="Write a poem")
        response = _make_text_response("Roses are red...")

        with patch("bulkgen.providers.mistral.Mistral") as mock_cls:
            mock_cls.return_value = _make_mistral_mock(response)

            provider = MistralProvider(api_key="test-key")
            await provider.generate(
                target_name="poem.txt",
                target_config=target_config,
                resolved_prompt="Write a poem",
                resolved_model=_model("mistral-large-latest"),
                project_dir=project_dir,
            )

        output = project_dir / "poem.txt"
        assert output.exists()
        assert output.read_text() == "Roses are red..."

    async def test_text_with_text_input(self, project_dir: Path) -> None:
        _ = (project_dir / "source.txt").write_text("Source material here")
        target_config = TargetConfig(prompt="Summarize", inputs=["source.txt"])
        response = _make_text_response("Summary: ...")

        with patch("bulkgen.providers.mistral.Mistral") as mock_cls:
            mock_client = _make_mistral_mock(response)
            mock_cls.return_value = mock_client

            provider = MistralProvider(api_key="test-key")
            await provider.generate(
                target_name="summary.md",
                target_config=target_config,
                resolved_prompt="Summarize",
                resolved_model=_model("mistral-large-latest"),
                project_dir=project_dir,
            )

            call_args = mock_client.chat.complete_async.call_args
            messages = call_args.kwargs["messages"]
            prompt_text = messages[0].content
            assert "--- Contents of source.txt ---" in prompt_text
            assert "Source material here" in prompt_text

    async def test_text_with_image_input(self, project_dir: Path) -> None:
        _ = (project_dir / "photo.png").write_bytes(b"\x89PNG")
        target_config = TargetConfig(prompt="Describe this image", inputs=["photo.png"])
        response = _make_text_response("A beautiful photo")

        with patch("bulkgen.providers.mistral.Mistral") as mock_cls:
            mock_client = _make_mistral_mock(response)
            mock_cls.return_value = mock_client

            provider = MistralProvider(api_key="test-key")
            await provider.generate(
                target_name="desc.txt",
                target_config=target_config,
                resolved_prompt="Describe this image",
                resolved_model=_model("mistral-large-latest"),
                project_dir=project_dir,
            )

            call_args = mock_client.chat.complete_async.call_args
            messages = call_args.kwargs["messages"]
            chunks = messages[0].content
            assert isinstance(chunks, list)
            assert chunks[0].text == "Describe this image"
            assert chunks[1].image_url.url.startswith("data:image/png;base64,")

    async def test_text_no_choices_raises(self, project_dir: Path) -> None:
        target_config = TargetConfig(prompt="x")
        response = MagicMock()
        response.choices = []

        with patch("bulkgen.providers.mistral.Mistral") as mock_cls:
            mock_cls.return_value = _make_mistral_mock(response)

            provider = MistralProvider(api_key="test-key")
            with pytest.raises(RuntimeError, match="no choices"):
                await provider.generate(
                    target_name="fail.txt",
                    target_config=target_config,
                    resolved_prompt="x",
                    resolved_model=_model("mistral-large-latest"),
                    project_dir=project_dir,
                )

    async def test_text_empty_content_raises(self, project_dir: Path) -> None:
        target_config = TargetConfig(prompt="x")
        response = _make_text_response(None)

        with patch("bulkgen.providers.mistral.Mistral") as mock_cls:
            mock_cls.return_value = _make_mistral_mock(response)

            provider = MistralProvider(api_key="test-key")
            with pytest.raises(RuntimeError, match="empty content"):
                await provider.generate(
                    target_name="fail.txt",
                    target_config=target_config,
                    resolved_prompt="x",
                    resolved_model=_model("mistral-large-latest"),
                    project_dir=project_dir,
                )

    async def test_text_with_multiple_inputs(self, project_dir: Path) -> None:
        _ = (project_dir / "a.txt").write_text("content A")
        _ = (project_dir / "b.txt").write_text("content B")
        _ = (project_dir / "c.png").write_bytes(b"\x89PNG")

        target_config = TargetConfig(
            prompt="Combine all", inputs=["a.txt", "b.txt", "c.png"]
        )
        response = _make_text_response("Combined")

        with patch("bulkgen.providers.mistral.Mistral") as mock_cls:
            mock_client = _make_mistral_mock(response)
            mock_cls.return_value = mock_client

            provider = MistralProvider(api_key="test-key")
            await provider.generate(
                target_name="out.md",
                target_config=target_config,
                resolved_prompt="Combine all",
                resolved_model=_model("mistral-large-latest"),
                project_dir=project_dir,
            )

            call_args = mock_client.chat.complete_async.call_args
            chunks = call_args.kwargs["messages"][0].content
            assert isinstance(chunks, list)
            # TextChunk for prompt, TextChunk for a.txt, TextChunk for b.txt,
            # ImageURLChunk for c.png
            assert chunks[0].text == "Combine all"
            assert "content A" in chunks[1].text
            assert "content B" in chunks[2].text
            assert chunks[3].image_url.url.startswith("data:image/png;base64,")

    async def test_text_with_reference_images(self, project_dir: Path) -> None:
        _ = (project_dir / "ref.png").write_bytes(b"\x89PNG")

        target_config = TargetConfig(
            prompt="Describe the style", reference_images=["ref.png"]
        )
        response = _make_text_response("A stylized image")

        with patch("bulkgen.providers.mistral.Mistral") as mock_cls:
            mock_client = _make_mistral_mock(response)
            mock_cls.return_value = mock_client

            provider = MistralProvider(api_key="test-key")
            await provider.generate(
                target_name="desc.txt",
                target_config=target_config,
                resolved_prompt="Describe the style",
                resolved_model=_model("mistral-large-latest"),
                project_dir=project_dir,
            )

            call_args = mock_client.chat.complete_async.call_args
            chunks = call_args.kwargs["messages"][0].content
            assert isinstance(chunks, list)
            assert chunks[0].text == "Describe the style"
            assert chunks[1].image_url.url.startswith("data:image/png;base64,")


def _make_openai_mock(b64_data: str) -> AsyncMock:
    """Return a mock AsyncOpenAI client that returns b64 image data."""
    image = MagicMock()
    image.b64_json = b64_data
    image.url = None

    response = MagicMock()
    response.data = [image]

    mock_client = AsyncMock()
    mock_client.images.generate = AsyncMock(return_value=response)
    mock_client.images.edit = AsyncMock(return_value=response)
    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
    mock_client.__aexit__ = AsyncMock(return_value=False)
    return mock_client


class TestOpenAIImageProvider:
    """Test OpenAIImageProvider with mocked OpenAI client."""

    @pytest.fixture
    def image_bytes(self) -> bytes:
        return b"\x89PNG\r\n\x1a\n" + b"\x00" * 100

    async def test_single_reference_image(
        self, project_dir: Path, image_bytes: bytes
    ) -> None:
        ref = project_dir / "ref.png"
        _ = ref.write_bytes(b"reference data")

        target_config = TargetConfig(prompt="Edit this", reference_images=["ref.png"])
        b64 = base64.b64encode(image_bytes).decode()
        mock_client = _make_openai_mock(b64)

        with patch("bulkgen.providers.openai_image.AsyncOpenAI") as mock_cls:
            mock_cls.return_value = mock_client

            provider = OpenAIImageProvider(api_key="test-key")
            await provider.generate(
                target_name="out.png",
                target_config=target_config,
                resolved_prompt="Edit this",
                resolved_model=_model("gpt-image-1"),
                project_dir=project_dir,
            )

            call_args = mock_client.images.edit.call_args
            # Single reference image should be passed as raw bytes
            assert call_args.kwargs["image"] == b"reference data"

        output = project_dir / "out.png"
        assert output.exists()
        assert output.read_bytes() == image_bytes

    async def test_multiple_reference_images(
        self, project_dir: Path, image_bytes: bytes
    ) -> None:
        ref1 = project_dir / "ref1.png"
        ref2 = project_dir / "ref2.png"
        _ = ref1.write_bytes(b"ref1 data")
        _ = ref2.write_bytes(b"ref2 data")

        target_config = TargetConfig(
            prompt="Combine", reference_images=["ref1.png", "ref2.png"]
        )
        b64 = base64.b64encode(image_bytes).decode()
        mock_client = _make_openai_mock(b64)

        with patch("bulkgen.providers.openai_image.AsyncOpenAI") as mock_cls:
            mock_cls.return_value = mock_client

            provider = OpenAIImageProvider(api_key="test-key")
            await provider.generate(
                target_name="out.png",
                target_config=target_config,
                resolved_prompt="Combine",
                resolved_model=_model("gpt-image-1"),
                project_dir=project_dir,
            )

            call_args = mock_client.images.edit.call_args
            # Multiple reference images should be passed as a list of bytes
            image_arg: list[bytes] = call_args.kwargs["image"]
            assert isinstance(image_arg, list)
            assert len(image_arg) == 2
            assert image_arg[0] == b"ref1 data"
            assert image_arg[1] == b"ref2 data"