import base64
import struct
import tempfile
import unittest
import uuid
from pathlib import Path
from unittest.mock import patch

from core.ark_client import ArkConfig
from tools.gpt_image import GptImageTool


def _png_stub(width: int = 1536, height: int = 864) -> bytes:
    return b"\x89PNG\r\n\x1a\n" + b"\x00\x00\x00\rIHDR" + struct.pack(">II", width, height)


class _FakeArkClient:
    json_call = None

    def __init__(self, *_args, **_kwargs):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        pass

    def post_json(self, endpoint, body, *, timeout_s=None):
        type(self).json_call = (endpoint, body, timeout_s)
        return {
            "data": [{"b64_json": base64.b64encode(_png_stub()).decode("ascii")}],
            "usage": {"output_tokens": 123},
        }


class GptImageToolTests(unittest.TestCase):
    def setUp(self):
        _FakeArkClient.json_call = None
        self.tmp = tempfile.TemporaryDirectory()
        self.root = Path(self.tmp.name)
        self.working_dir = self.root / "task"
        self.working_dir.mkdir()
        self.cfg = {
            "model_id": "gpt-image-2",
            "endpoint": "/images/generations",
            "default_size": "auto",
            "default_quality": "auto",
            "request_timeout_s": 300,
            "price_cny_per_image": 0,
        }
        self.tool = GptImageTool(
            gw_cfg=ArkConfig(api_key="test", base_url="https://example.test/v1", raw={}),
            image_variant_cfg=self.cfg,
            variant_key="gpt_image",
            working_dir=self.working_dir,
            task_id=uuid.uuid4(),
            user_id=uuid.uuid4(),
            base_dir=self.working_dir,
            user_root=self.root,
            daily_limit=0,
        )

    def tearDown(self):
        self.tmp.cleanup()

    def _execute(self, **kwargs):
        with (
            patch("tools.gpt_image.ArkClient", _FakeArkClient),
            patch("tools.gpt_image.quota_gate", return_value=""),
            patch("tools.gpt_image.record_usage_safe"),
        ):
            return self.tool.execute(**kwargs)

    def test_text_to_image_forwards_size_and_quality(self):
        result = self._execute(
            prompt="draw a materials lab",
            size="1536x864",
            quality="high",
        )

        self.assertTrue(result.startswith("[gpt_image]"))
        endpoint, body, _timeout = _FakeArkClient.json_call
        self.assertEqual(endpoint, "/images/generations")
        self.assertEqual(body["size"], "1536x864")
        self.assertEqual(body["quality"], "high")
        self.assertIn("size=1536x864", result)
        self.assertIn("quality=high", result)

    def test_size_validation(self):
        self.assertEqual(self.tool._normalize_size("auto"), ("auto", ""))
        self.assertEqual(self.tool._normalize_size("1536×864"), ("1536x864", ""))
        for value in ("1000x1000", "4096x1024", "3072x512", "640x640", "bad"):
            with self.subTest(value=value):
                _size, error = self.tool._normalize_size(value)
                self.assertTrue(error.startswith("[Error]"))

    def test_quality_validation(self):
        result = self._execute(prompt="draw", quality="ultra")
        self.assertEqual(result, "[Error] quality 必须是 auto / low / medium / high")
        self.assertIsNone(_FakeArkClient.json_call)


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