from io import BytesIO
import unittest
from unittest.mock import AsyncMock, patch

from PIL import Image
from pydantic import ValidationError

from models import AssistantConfig
from services.test_assets import read_test_image, store_test_image
from services.test_runs.text_runner import _transcript_content
from services.vision import build_image_user_message
from test_schemas import FixedInputTurn, TestImageRef


def _jpeg_bytes() -> bytes:
    buffer = BytesIO()
    Image.new("RGB", (12, 8), (220, 40, 40)).save(buffer, format="JPEG")
    return buffer.getvalue()


def _image_ref() -> TestImageRef:
    return TestImageRef(
        asset_id="ta_fixture",
        filename="fixture.jpg",
        width=12,
        height=8,
        size_bytes=128,
    )


class TestImageSchemaTests(unittest.TestCase):
    def test_allows_image_only_turn(self):
        turn = FixedInputTurn(id="turn_1", image=_image_ref())

        self.assertEqual(turn.user_input, "")
        self.assertEqual(turn.image.asset_id, "ta_fixture")

    def test_rejects_turn_without_text_or_image(self):
        with self.assertRaisesRegex(ValidationError, "至少填写一项"):
            FixedInputTurn(id="turn_1")


class PersistentTestAssetTests(unittest.TestCase):
    def test_normalizes_and_stores_under_test_prefix(self):
        with patch("services.test_assets.put_object") as put_object:
            stored = store_test_image("ta_fixture", _jpeg_bytes())

        self.assertEqual(stored.storage_key, "test-assets/ta_fixture.jpg")
        self.assertEqual((stored.width, stored.height), (12, 8))
        self.assertEqual(put_object.call_args.args[0], stored.storage_key)
        self.assertEqual(put_object.call_args.args[2], "image/jpeg")

    def test_reads_only_test_asset_keys(self):
        with patch("services.test_assets.get_object", return_value=b"jpeg"):
            self.assertEqual(read_test_image("test-assets/ta_fixture.jpg"), b"jpeg")
        with self.assertRaisesRegex(ValueError, "存储位置无效"):
            read_test_image("conversation-inputs/unsafe.jpg")


class ImageUserMessageTests(unittest.IsolatedAsyncioTestCase):
    async def test_native_model_receives_multimodal_message(self):
        message = await build_image_user_message(
            AssistantConfig(
                type="prompt",
                vision_enabled=True,
                llm_support_image_input=True,
            ),
            _jpeg_bytes(),
            "这是什么颜色？",
        )

        self.assertEqual(message["role"], "user")
        self.assertEqual(message["content"][0]["text"], "这是什么颜色？")
        self.assertTrue(
            message["content"][1]["image_url"]["url"].startswith(
                "data:image/jpeg;base64,"
            )
        )

    async def test_auxiliary_model_observation_becomes_text_context(self):
        with patch(
            "services.vision.analyze_image_with_vision_model",
            new=AsyncMock(return_value="画面中是一个红色方块。"),
        ):
            message = await build_image_user_message(
                AssistantConfig(
                    type="prompt",
                    vision_enabled=True,
                    vision_model_resource_id="vision_model",
                    vision_llm_support_image_input=True,
                ),
                _jpeg_bytes(),
                "画面里有什么？",
            )

        self.assertIsInstance(message["content"], str)
        self.assertIn("红色方块", message["content"])

    def test_transcript_does_not_persist_image_data(self):
        content = _transcript_content(
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "请看图"},
                    {
                        "type": "image_url",
                        "image_url": {"url": "data:image/jpeg;base64,secret"},
                    },
                ],
            }
        )

        self.assertEqual(content, "请看图\n[用户提交了一张测试图片]")
        self.assertNotIn("base64", content)


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