from __future__ import annotations

import unittest
from types import SimpleNamespace
from unittest import mock

from services.protocol import conversation


class ImageParallelGenerationTests(unittest.TestCase):
    def test_multi_image_request_uses_independent_pool_jobs_when_enabled(self) -> None:
        calls: list[tuple[int, int, int | None]] = []

        def fake_serial(request, *, index_offset=0, total_count=None, **_kwargs):
            calls.append((request.n, index_offset, total_count))
            yield conversation.ImageOutput(
                kind="result",
                model=request.model,
                index=index_offset + 1,
                total=int(total_count or request.n),
                data=[{"url": f"https://example.test/{index_offset + 1}.png"}],
            )

        with (
            mock.patch.object(conversation, "config", SimpleNamespace(image_parallel_generation=True)),
            mock.patch.object(conversation, "_stream_image_outputs_with_pool", side_effect=fake_serial),
        ):
            outputs = list(conversation.stream_image_outputs_with_pool(
                conversation.ConversationRequest(model="gpt-image-2", prompt="cat", n=3)
            ))

        self.assertEqual(sorted(calls), [(1, 0, 3), (1, 1, 3), (1, 2, 3)])
        self.assertEqual(sorted(output.index for output in outputs), [1, 2, 3])
        self.assertTrue(all(output.total == 3 for output in outputs))

    def test_multi_image_request_stays_serial_when_disabled(self) -> None:
        calls: list[int] = []

        def fake_serial(request, **_kwargs):
            calls.append(request.n)
            for index in range(1, request.n + 1):
                yield conversation.ImageOutput(
                    kind="result",
                    model=request.model,
                    index=index,
                    total=request.n,
                    data=[{"url": f"https://example.test/{index}.png"}],
                )

        with (
            mock.patch.object(conversation, "config", SimpleNamespace(image_parallel_generation=False)),
            mock.patch.object(conversation, "_stream_image_outputs_with_pool", side_effect=fake_serial),
        ):
            outputs = list(conversation.stream_image_outputs_with_pool(
                conversation.ConversationRequest(model="gpt-image-2", prompt="cat", n=2)
            ))

        self.assertEqual(calls, [2])
        self.assertEqual([output.index for output in outputs], [1, 2])


if __name__ == "__main__":
    unittest.main()