from __future__ import annotations

import base64
import io
import unittest
from unittest import mock

from PIL import Image

from services import image_annotation_service


def _image_data_url(width: int = 200, height: int = 100) -> str:
    output = io.BytesIO()
    Image.new("RGB", (width, height), (245, 240, 235)).save(output, format="PNG")
    return "data:image/png;base64," + base64.b64encode(output.getvalue()).decode("ascii")


class ImageAnnotationServiceTests(unittest.TestCase):
    def test_identify_region_sends_original_and_crop_to_configured_agent(self) -> None:
        captured: dict[str, object] = {}

        def fake_agent(identity, system_prompt, user_prompt, image_data_urls, **kwargs):
            captured.update(
                identity=identity,
                system_prompt=system_prompt,
                user_prompt=user_prompt,
                images=image_data_urls,
                excluded_tokens=kwargs.get("excluded_tokens"),
            )
            return (
                '{"label":"矿泉水瓶","object_type":"object","visible_text":"PURE",'
                '"description":"透明塑料瓶","confidence":0.93}',
                "vision-agent",
            )

        with mock.patch.object(image_annotation_service, "call_product_visual_agent", side_effect=fake_agent):
            result = image_annotation_service.identify_image_region(
                {"id": "user-1"},
                _image_data_url(),
                {"x": 10, "y": 20, "width": 30, "height": 40},
            )

        self.assertEqual(result["label"], "矿泉水瓶")
        self.assertEqual(result["model"], "vision-agent")
        self.assertEqual(result["image_width"], 200)
        self.assertEqual(result["image_height"], 100)
        self.assertEqual(result["bounds"], {"x": 10.0, "y": 20.0, "width": 30.0, "height": 40.0})
        images = captured["images"]
        self.assertEqual(len(images), 2)
        self.assertTrue(images[0].startswith("data:image/png;base64,"))
        self.assertTrue(images[1].startswith("data:image/jpeg;base64,"))

    def test_identify_region_retries_malformed_agent_output(self) -> None:
        responses = [
            ("", "vision-agent"),
            ("not json", "vision-agent"),
            ('{"label":"人物的头发","object_type":"person_part","confidence":2}', "vision-agent"),
        ]
        with mock.patch.object(image_annotation_service, "call_product_visual_agent", side_effect=responses) as call:
            result = image_annotation_service.identify_image_region(
                {"id": "user-1"},
                _image_data_url(),
                {"x": 25, "y": 10, "width": 20, "height": 25},
            )

        self.assertEqual(call.call_count, 3)
        self.assertEqual(result["label"], "人物的头发")
        self.assertEqual(result["confidence"], 1.0)


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