"""Model-independent image decode parity for native Rust MM.

Covers ``decode_rgb`` in ``rust/sglang-mm/src/common/mod.rs`` (via the
``_core.common.image_decode_rgb`` binding), against PIL's
``Image.open(...).convert("RGB")``.
"""

import io
import sys
import unittest
from pathlib import Path

import numpy as np
from PIL import Image

from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from _mm_rust_utils import load_core  # noqa: E402

register_cpu_ci(est_time=5, suite="base-a-test-cpu")

CORE = load_core()
DECODE = CORE and CORE.common.image_decode_rgb


def encode(image, fmt, **kwargs):
    buffer = io.BytesIO()
    image.save(buffer, format=fmt, **kwargs)
    return buffer.getvalue()


@unittest.skipUnless(DECODE, "sglang-mm decode binding not built")
class TestRustImageDecode(CustomTestCase):
    def assert_matches_pil(self, data, tolerance=0):
        height, width, pixels = DECODE(data)
        expected = np.asarray(Image.open(io.BytesIO(data)).convert("RGB"))
        self.assertEqual((height, width), expected.shape[:2])
        actual = np.asarray(pixels).reshape(expected.shape)
        diff = np.abs(actual.astype(int) - expected.astype(int))
        self.assertLessEqual(diff.max(), tolerance)

    def test_png_modes_match_pil(self):
        rgb = np.random.default_rng(1).integers(0, 256, (19, 23, 3), dtype=np.uint8)
        cases = [
            ("RGB", Image.fromarray(rgb)),
            ("L", Image.fromarray(rgb[..., 0])),
            ("RGBA", Image.fromarray(np.dstack((rgb, rgb[..., 0])))),
            ("P", Image.fromarray(rgb).quantize(colors=16)),
        ]
        for mode, image in cases:
            with self.subTest(mode=mode):
                self.assert_matches_pil(encode(image, "PNG"))

    def test_jpeg_modes_match_with_decoder_tolerance(self):
        rgb = np.random.default_rng(2).integers(0, 256, (31, 29, 3), dtype=np.uint8)
        image = Image.fromarray(rgb)
        exif = image.getexif()
        exif[274] = 6  # EXIF orientation: neither decoder applies it
        cases = [
            ("RGB", encode(image, "JPEG")),
            ("L", encode(Image.fromarray(rgb[..., 0]), "JPEG")),
            ("CMYK", encode(image.convert("CMYK"), "JPEG")),
            ("EXIF-rotated", encode(image, "JPEG", exif=exif)),
        ]
        for mode, data in cases:
            with self.subTest(mode=mode):
                self.assert_matches_pil(data, tolerance=3)

    def test_lossless_formats_match_pil_exactly(self):
        """GIF/BMP/lossless-WebP joined the native decoder with the pure-Rust
        webp/gif/bmp enablement. Their reconstruction is exact by format spec,
        so parity with PIL is pinned at zero tolerance like PNG."""
        rgb = np.random.default_rng(4).integers(0, 256, (17, 21, 3), dtype=np.uint8)
        image = Image.fromarray(rgb)
        cases = {
            "gif": encode(image.quantize(colors=64), "GIF"),
            "bmp": encode(image, "BMP"),
            "webp-lossless": encode(image, "WEBP", lossless=True),
        }
        for name, data in cases.items():
            with self.subTest(fmt=name):
                self.assert_matches_pil(data)

    def test_unsupported_inputs_fail(self):
        gray16 = Image.fromarray(np.zeros((9, 9), dtype=np.uint16))
        cases = {
            "corrupt": b"not an image",
            # >8-bit depths must error (the request is then rejected — there
            # is no Python fallback) — never silently diverge from PIL's
            # clipping.
            "png16": encode(gray16, "PNG"),
        }
        for name, data in cases.items():
            with self.subTest(input=name):
                with self.assertRaises(ValueError):
                    DECODE(data)


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