from __future__ import annotations

import os
import unittest
from unittest.mock import patch

import gpt_image2_generate as image_demo
import gpt_image2_generate_multi as image_multi_demo


class ImageDemoHelpersTest(unittest.TestCase):
    def test_openai_ratio_size_mapping(self) -> None:
        self.assertEqual(image_demo.resolve_api_size("16:9", "openai"), "1536x864")
        self.assertEqual(image_demo.resolve_api_size("3840x2160", "openai"), "3840x2160")

    def test_min_output_edge(self) -> None:
        self.assertEqual(image_demo.resolve_min_output_edge("4096x2304", None, "apicodex"), 4096)
        self.assertIsNone(image_demo.resolve_min_output_edge("1024x1024", None, "openai"))


class ImageMultiDemoHelpersTest(unittest.TestCase):
    def test_collect_reference_sources_from_json_array(self) -> None:
        with patch.dict(os.environ, {"REFERENCE_IMAGES": '["./a.png", "./b.png", "./a.png"]'}, clear=True):
            self.assertEqual(image_multi_demo.collect_reference_sources(), ["./a.png", "./b.png"])

    def test_build_edit_files_field_names(self) -> None:
        ref1 = image_demo.ImageInput("a.png", "image/png", b"a")
        ref2 = image_demo.ImageInput("b.png", "image/png", b"b")

        openai_single = image_multi_demo.build_edit_files("openai", [ref1], None)
        openai_multi = image_multi_demo.build_edit_files("openai", [ref1, ref2], None)
        apicodex_multi = image_multi_demo.build_edit_files("apicodex", [ref1, ref2], None)

        self.assertEqual([name for name, _ in openai_single], ["image"])
        self.assertEqual([name for name, _ in openai_multi], ["image[]", "image[]"])
        self.assertEqual([name for name, _ in apicodex_multi], ["image[]", "image[]"])

    def test_output_stem_env_is_sanitized(self) -> None:
        with patch.dict(os.environ, {"IMAGE_OUTPUT_STEM": "../bad stem!"}, clear=True):
            self.assertEqual(image_multi_demo.resolve_output_stem(), "bad_stem")

    def test_default_n_matches_reference_count(self) -> None:
        env = {
            "IMAGE_PROMPT": "hello",
            "IMAGE_SIZE": "1:1",
        }
        with patch.dict(os.environ, env, clear=True):
            payload, model, size, request_mode, mask_source, background, input_fidelity = image_multi_demo.build_payload(
                "http://localhost",
                ["./a.png", "./b.png"],
            )

        self.assertEqual(payload["n"], 2)
        self.assertEqual(request_mode, "edit")
        self.assertIsNone(mask_source)
        self.assertIsNone(background)
        self.assertIsNone(input_fidelity)
        self.assertEqual(model, image_demo.DEFAULT_MODEL)
        self.assertEqual(size, "1:1")


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