import sys
import tempfile
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_da3


class FakeSession:
    def get_inputs(self):
        return [SimpleNamespace(name="images"), SimpleNamespace(name="strategy")]

    def run(self, output_names, feeds):
        self.feeds = feeds
        _, views, _, height, width = feeds["images"].shape
        depth = np.ones((1, views, height, width), dtype=np.float32)
        confidence = np.full_like(depth, 0.75)
        pose = np.zeros((1, views, 9), dtype=np.float32)
        pose[..., 6] = 1.0
        pose[..., 7:] = np.pi / 2.0
        return depth, confidence, pose


class Da3OnnxTest(unittest.TestCase):
    def test_target_size_uses_patch_multiple(self):
        self.assertEqual(
            onnx_da3.target_size((800, 600), 504, "upper_bound_resize"),
            (504, 378),
        )

    def test_infer_returns_camera_and_image_arrays(self):
        with tempfile.TemporaryDirectory() as directory:
            image_path = Path(directory) / "input.png"
            Image.new("RGB", (80, 60), (10, 20, 30)).save(image_path)
            session = FakeSession()

            result = onnx_da3.infer(
                session,
                [image_path],
                process_res=56,
                ref_view_strategy="middle",
            )

        self.assertEqual(result.depth.shape, (1, 42, 56))
        self.assertEqual(result.conf.shape, result.depth.shape)
        self.assertEqual(result.processed_images.shape, (1, 42, 56, 3))
        self.assertEqual(result.extrinsics.shape, (1, 3, 4))
        self.assertEqual(result.intrinsics.shape, (1, 3, 3))
        self.assertEqual(int(session.feeds["strategy"][0]), 1)
        np.testing.assert_allclose(result.extrinsics[0, :3, :3], np.eye(3))

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