"""Model-independent parity tests for Rust multimodal source loading.

Covers ``fetch_bytes`` in ``rust/sglang-mm/src/common/fetch.rs`` (via the
``_core.common.fetch_bytes`` binding), against the Python reference
``sglang.srt.utils.common.get_image_bytes``.
"""

import base64
import http.server
import sys
import tempfile
import threading
import unittest
from pathlib import Path

from sglang.srt.utils.common import get_image_bytes
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()
FETCH = CORE and CORE.common.fetch_bytes


@unittest.skipUnless(FETCH, "sglang-mm fetch binding not built")
class TestRustMediaSourceLoading(CustomTestCase):
    DATA = b"native-mm-source"

    def test_inline_sources(self):
        encoded = base64.b64encode(self.DATA).decode()
        for source in (encoded, f"data:application/octet-stream;base64,{encoded}"):
            with self.subTest(source=source[:8]):
                self.assertEqual(bytes(FETCH(source)), get_image_bytes(source))

    def test_file_sources(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "image.png"
            path.write_bytes(self.DATA)
            self.assertEqual(bytes(FETCH(str(path))), get_image_bytes(str(path)))
            # `file://` is asserted against the payload, not the Python helper,
            # on purpose: `get_image_bytes` passes the un-stripped URL straight
            # to `open()`, so the reference raises here. The native path strips
            # the scheme and succeeds — a deliberate divergence, not an omission.
            with self.assertRaises(OSError):
                get_image_bytes(path.as_uri())
            self.assertEqual(bytes(FETCH(path.as_uri())), self.DATA)

    def test_http_source(self):
        data = self.DATA

        class Handler(http.server.BaseHTTPRequestHandler):
            def do_GET(self):
                self.send_response(200)
                self.end_headers()
                self.wfile.write(data)

            def log_message(self, *_):
                pass

        server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        try:
            url = f"http://127.0.0.1:{server.server_port}/image"
            self.assertEqual(bytes(FETCH(url)), get_image_bytes(url))
        finally:
            server.shutdown()
            server.server_close()
            thread.join()

    def test_invalid_sources_fail(self):
        for source in ("not base64!", "/definitely/missing/image.png"):
            with self.subTest(source=source):
                with self.assertRaises(ValueError):
                    FETCH(source)


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