"""仅扩图的像素保护与单张 AI 编辑提示词回归测试。"""
from __future__ import annotations

import asyncio
import base64
import io
import sys
import unittest
from pathlib import Path

from PIL import Image

SERVER_DIR = Path(__file__).resolve().parents[1]
if str(SERVER_DIR) not in sys.path:
    sys.path.insert(0, str(SERVER_DIR))

from services.generator import prepare_outpaint_canvas
from services.prompts import build_prompt


def _data_url(image: Image.Image) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode()


class OutpaintTests(unittest.TestCase):
    def test_outpaint_restores_original_region_after_provider_response(self) -> None:
        original = Image.new("RGB", (400, 400), (220, 20, 60))
        canvas = asyncio.run(prepare_outpaint_canvas(_data_url(original), "1152x1536"))

        self.assertEqual(canvas.size, (1152, 1536))
        self.assertEqual(canvas.box, (0, 192, 1152, 1344))

        # 模拟供应商把整张图都重绘成蓝色；服务端合成后中心原图必须仍为红色。
        generated = Image.new("RGB", canvas.size, (0, 0, 255))
        restored = Image.open(io.BytesIO(canvas.restore_original(_png_bytes(generated)))).convert("RGB")
        self.assertEqual(restored.getpixel((576, 768)), (220, 20, 60))
        self.assertEqual(restored.getpixel((576, 80)), (0, 0, 255))

    def test_expanded_margin_is_transparent_in_mask(self) -> None:
        original = Image.new("RGB", (400, 400), (1, 2, 3))
        canvas = asyncio.run(prepare_outpaint_canvas(_data_url(original), "1152x1536"))
        mask_bytes = base64.b64decode(canvas.mask_data_url.split(",", 1)[1])
        mask = Image.open(io.BytesIO(mask_bytes)).convert("RGBA")
        self.assertEqual(mask.getpixel((576, 80))[3], 0)
        self.assertEqual(mask.getpixel((576, 768))[3], 255)


def _png_bytes(image: Image.Image) -> bytes:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")
    return buffer.getvalue()


class ExclusiveEditPromptTests(unittest.TestCase):
    def test_user_edit_requirements_do_not_inherit_marketing_or_light_studio_prompt(self) -> None:
        prompt = build_prompt(
            "pcore", "gpt-image-2", "custom", {"title": "product"}, 5, "ru",
            requirements="背景必须是纯黑色，其他不要改变",
            requirements_exclusive=True,
        )
        self.assertIn("背景必须是纯黑色", prompt)
        self.assertNotIn("e-commerce marketing image", prompt)
        self.assertNotIn("plain light background", prompt)
        self.assertNotIn("Render concise Russian marketing text", prompt)


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