from __future__ import annotations

import base64
import contextlib
import importlib.util
import io
import os
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest import mock


ROOT = Path(__file__).resolve().parents[1]
SCRIPT = ROOT / "scripts" / "eapil_gpt2_image.py"
SKILL = ROOT / "SKILL.md"


def load_module():
    spec = importlib.util.spec_from_file_location("eapil_gpt2_image", SCRIPT)
    module = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(module)
    return module


class EapilGpt2ImageTests(unittest.TestCase):
    def setUp(self):
        self.module = load_module()

    def test_api_key_prefers_cli_key_then_eapil_then_openai(self):
        with mock.patch.dict(
            os.environ,
            {"EAPIL_API_KEY": "eapil", "OPENAI_API_KEY": "openai"},
            clear=True,
        ):
            self.assertEqual(self.module.resolve_api_key("direct"), "direct")
            self.assertEqual(self.module.resolve_api_key(None), "eapil")
        with mock.patch.dict(os.environ, {"OPENAI_API_KEY": "openai"}, clear=True):
            self.assertEqual(self.module.resolve_api_key(None), "openai")

    def test_missing_api_key_mentions_all_supported_sources(self):
        with mock.patch.dict(os.environ, {}, clear=True):
            stderr = io.StringIO()
            with contextlib.redirect_stderr(stderr), self.assertRaises(SystemExit):
                self.module.resolve_api_key(None)
            message = stderr.getvalue()
            self.assertIn("--api-key", message)
            self.assertIn("EAPIL_API_KEY", message)
            self.assertIn("OPENAI_API_KEY", message)

    def test_generation_payload_omits_absent_optional_values_and_keeps_supplied_values(self):
        captured = {}

        def fake_request_json(path, payload, key):
            captured["path"] = path
            captured["payload"] = payload
            captured["key"] = key
            return {
                "data": [{"b64_json": base64.b64encode(b"png").decode("ascii")}]
            }

        self.module.request_json = fake_request_json
        self.module.save_response_images = lambda response, path, n: None
        args = SimpleNamespace(
            prompt="test prompt",
            size="1024x1024",
            output=None,
            api_key="direct",
            quality="medium",
            n=2,
            background="opaque",
            moderation="low",
            output_format="webp",
            output_compression=80,
            user="tester",
        )

        self.module.command_generate(args)

        self.assertEqual(captured["path"], "/images/generations")
        self.assertEqual(captured["key"], "direct")
        self.assertEqual(captured["payload"]["quality"], "medium")
        self.assertEqual(captured["payload"]["n"], 2)
        self.assertEqual(captured["payload"]["background"], "opaque")
        self.assertEqual(captured["payload"]["moderation"], "low")
        self.assertEqual(captured["payload"]["output_format"], "webp")
        self.assertEqual(captured["payload"]["output_compression"], 80)
        self.assertEqual(captured["payload"]["user"], "tester")

    def test_save_response_images_suffixes_multiple_outputs(self):
        response = {
            "data": [
                {"b64_json": base64.b64encode(b"one").decode("ascii")},
                {"b64_json": base64.b64encode(b"two").decode("ascii")},
            ]
        }
        with tempfile.TemporaryDirectory() as tmp:
            target = Path(tmp) / "image.png"
            with contextlib.redirect_stdout(io.StringIO()):
                outputs = self.module.save_response_images(response, target, 2)
            self.assertEqual([p.name for p in outputs], ["image_0.png", "image_1.png"])
            self.assertEqual(outputs[0].read_bytes(), b"one")
            self.assertEqual(outputs[1].read_bytes(), b"two")

    def test_parser_accepts_repeated_reference_images(self):
        args = self.module.parser().parse_args(
            [
                "edit",
                "--prompt",
                "combine",
                "--image",
                "one.png",
                "--image",
                "two.png",
            ]
        )
        self.assertEqual(args.image, ["one.png", "two.png"])

    def test_skill_documents_prompt_enhancement_gate_and_key_fallbacks(self):
        text = SKILL.read_text(encoding="utf-8")
        self.assertIn("OPENAI_API_KEY", text)
        self.assertIn("--api-key", text)
        self.assertIn("prompt enhancement", text.lower())
        self.assertIn("confirm", text.lower())
        self.assertIn("references/gallery.md", text)
        self.assertIn("references/craft.md", text)
        self.assertIn("API execution prompt", text)
        self.assertIn("中文对照", text)
        self.assertIn("In-image text", text)
        self.assertIn("review only", text.lower())
        self.assertIn("user-requested language", text.lower())


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