import sys
import unittest
from pathlib import Path
from types import SimpleNamespace

import numpy as np
from PIL import Image


SOURCE_DIR = (
    Path(__file__).resolve().parents[1]
    / "src"
    / "anyimage"
    / "server"
    / "models"
)
sys.path.insert(0, str(SOURCE_DIR))

import onnx_realesrgan


class FakeSession:
    def get_inputs(self):
        return [SimpleNamespace(name="input")]

    def get_outputs(self):
        return [SimpleNamespace(name="output")]

    def run(self, output_names, feeds):
        self.output_names = output_names
        self.feeds = feeds
        source = feeds["input"]
        return [np.repeat(np.repeat(source, 4, axis=2), 4, axis=3)]


class RealesrganOnnxTest(unittest.TestCase):
    def test_preprocess_matches_dynamic_rgb_layout(self):
        image = Image.new("RGB", (7, 5), (64, 128, 192))

        result = onnx_realesrgan.preprocess(image)

        self.assertEqual(result.shape, (1, 3, 5, 7))
        self.assertEqual(result.dtype, np.float32)
        self.assertTrue(result.flags.c_contiguous)
        np.testing.assert_allclose(
            result[0, :, 0, 0],
            np.array((64, 128, 192), dtype=np.float32) / 255.0,
        )

    def test_infer_returns_direct_four_x_rgb_prediction(self):
        session = FakeSession()
        image = Image.new("RGB", (7, 5), (10, 20, 30))

        result = onnx_realesrgan.infer(session, image)

        self.assertEqual(result.mode, "RGB")
        self.assertEqual(result.size, (28, 20))
        self.assertEqual(session.output_names, ["output"])
        self.assertEqual(session.feeds["input"].shape, (1, 3, 5, 7))
        self.assertEqual(result.getpixel((0, 0)), (10, 20, 30))

    def test_infer_rejects_non_four_x_output(self):
        class InvalidSession(FakeSession):
            def run(self, _output_names, feeds):
                return [feeds["input"]]

        with self.assertRaisesRegex(RuntimeError, "unexpected output shape"):
            onnx_realesrgan.infer(
                InvalidSession(),
                Image.new("RGB", (7, 5)),
            )


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