import numpy as np
import pyzed.sl as sl
from unittest.mock import MagicMock
from aruco.svo_sync import SVOReader


def test_retrieve_depth_unit_guard():
    # Setup SVOReader with depth enabled
    reader = SVOReader([], depth_mode=sl.DEPTH_MODE.ULTRA)

    # Mock Camera
    mock_cam = MagicMock(spec=sl.Camera)

    # Mock depth data (e.g., 2.0 meters)
    depth_data = np.full((100, 100), 2.0, dtype=np.float32)
    mock_mat = MagicMock(spec=sl.Mat)
    mock_mat.get_data.return_value = depth_data

    # Mock retrieve_measure to "fill" the mat
    mock_cam.retrieve_measure.return_value = sl.ERROR_CODE.SUCCESS

    # Case 1: Units are METER -> Should NOT divide by 1000
    mock_init_params_meter = MagicMock(spec=sl.InitParameters)
    mock_init_params_meter.coordinate_units = sl.UNIT.METER
    mock_cam.get_init_parameters.return_value = mock_init_params_meter

    # We need to patch sl.Mat in the test or just rely on the fact that
    # _retrieve_depth creates a new sl.Mat() and calls get_data() on it.
    # Since we can't easily mock the sl.Mat() call inside the method without patching,
    # let's use a slightly different approach: mock the sl.Mat class itself.

    with MagicMock() as mock_mat_class:
        from aruco import svo_sync

        original_mat = svo_sync.sl.Mat
        svo_sync.sl.Mat = mock_mat_class
        mock_mat_instance = mock_mat_class.return_value
        mock_mat_instance.get_data.return_value = depth_data

        # Test METER path
        depth_meter = reader._retrieve_depth(mock_cam)
        assert depth_meter is not None
        assert np.allclose(depth_meter, 2.0)

        # Case 2: Units are MILLIMETER -> Should divide by 1000
        mock_init_params_mm = MagicMock(spec=sl.InitParameters)
        mock_init_params_mm.coordinate_units = sl.UNIT.MILLIMETER
        mock_cam.get_init_parameters.return_value = mock_init_params_mm

        depth_mm = reader._retrieve_depth(mock_cam)
        assert depth_mm is not None
        assert np.allclose(depth_mm, 0.002)

        # Restore original sl.Mat
        svo_sync.sl.Mat = original_mat


if __name__ == "__main__":
    test_retrieve_depth_unit_guard()