from __future__ import annotations

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

from PIL import Image

from services.input_assets import consume_input_image, store_input_image


def png_bytes(size: tuple[int, int] = (32, 24)) -> bytes:
    output = BytesIO()
    Image.new("RGBA", size, (20, 80, 160, 180)).save(output, format="PNG")
    return output.getvalue()


class InputAssetTests(unittest.TestCase):
    def test_store_normalizes_and_consume_removes_temporary_object(self):
        with (
            patch("services.input_assets.put_object") as put_object,
            patch("services.input_assets.get_object") as get_object,
            patch("services.input_assets.delete_object") as delete_object,
        ):
            stored = store_input_image(png_bytes())

            key, normalized, mime_type = put_object.call_args.args
            self.assertTrue(key.startswith("conversation-inputs/"))
            self.assertEqual(mime_type, "image/jpeg")
            self.assertTrue(normalized.startswith(b"\xff\xd8"))
            self.assertEqual((stored.width, stored.height), (32, 24))
            self.assertEqual(stored.size_bytes, len(normalized))

            get_object.return_value = normalized
            consumed = consume_input_image(stored.token)

        self.assertEqual(consumed, normalized)
        get_object.assert_called_once_with(key)
        delete_object.assert_called_once_with(key)

    def test_tampered_token_is_rejected_before_storage_read(self):
        with patch("services.input_assets.put_object"):
            stored = store_input_image(png_bytes())
        tampered = f"{stored.token}x"

        with patch("services.input_assets.get_object") as get_object:
            with self.assertRaisesRegex(ValueError, "签名无效"):
                consume_input_image(tampered)

        get_object.assert_not_called()

    def test_non_image_is_rejected(self):
        with self.assertRaisesRegex(ValueError, "可识别的图片"):
            store_input_image(b"not an image")


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