import pytest
import numpy as np
from unittest.mock import MagicMock, patch
import sys
from pathlib import Path

# Add py_workspace to path
sys.path.append(str(Path(__file__).parent.parent))

from calibrate_extrinsics import apply_depth_verify_refine_postprocess


@pytest.fixture
def mock_dependencies():
    with (
        patch("calibrate_extrinsics.verify_extrinsics_with_depth") as mock_verify,
        patch("calibrate_extrinsics.refine_extrinsics_with_depth") as mock_refine,
        patch("calibrate_extrinsics.click.echo") as mock_echo,
        patch("calibrate_extrinsics.save_depth_data") as mock_save_depth,
    ):
        # Setup mock return values
        mock_verify_res = MagicMock()
        mock_verify_res.rmse = 0.05
        mock_verify_res.mean_abs = 0.04
        mock_verify_res.median = 0.03
        mock_verify_res.depth_normalized_rmse = 0.02
        mock_verify_res.n_valid = 100
        mock_verify_res.n_total = 120
        mock_verify_res.residuals = []
        mock_verify.return_value = mock_verify_res

        mock_refine.return_value = (np.eye(4), {"success": True})

        yield mock_verify, mock_refine, mock_echo, mock_save_depth


def test_save_depth_data_integration(mock_dependencies):
    """
    Test that save_depth_data is called with correct data when save_depth_path is provided.
    """
    mock_verify, _, _, mock_save_depth = mock_dependencies

    serial = "123456"
    serial_int = int(serial)
    results = {serial: {"pose": "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"}}

    # Create frames
    # Frame 1
    f1 = MagicMock()
    f1.depth_map = np.ones((10, 20)) * 2.0  # H=10, W=20
    f1.confidence_map = np.zeros((10, 20))

    # Frame 2
    f2 = MagicMock()
    f2.depth_map = np.ones((10, 20)) * 2.2
    f2.confidence_map = np.zeros((10, 20))

    vfs = [
        {
            "frame": f1,
            "ids": np.array([[1]]),
            "corners": np.zeros((1, 4, 2)),
            "score": 100,
            "frame_index": 10,
        },
        {
            "frame": f2,
            "ids": np.array([[1]]),
            "corners": np.zeros((1, 4, 2)),
            "score": 90,
            "frame_index": 20,
        },
    ]

    verification_frames = {serial_int: vfs}
    marker_geometry = {1: np.zeros((4, 3))}
    camera_matrices = {serial_int: np.eye(3)}

    save_path = "output_depth.h5"

    # Run with save_depth_path
    apply_depth_verify_refine_postprocess(
        results=results,
        verification_frames=verification_frames,
        marker_geometry=marker_geometry,
        camera_matrices=camera_matrices,
        verify_depth=True,
        refine_depth=False,
        use_confidence_weights=False,
        depth_confidence_threshold=50,
        depth_pool_size=2,
        save_depth_path=save_path,
    )

    # Verify save_depth_data was called
    mock_save_depth.assert_called_once()

    # Check arguments
    call_args = mock_save_depth.call_args
    assert call_args[0][0] == save_path

    camera_data = call_args[0][1]
    assert serial in camera_data

    cam_data = camera_data[serial]
    assert "intrinsics" in cam_data
    assert "resolution" in cam_data
    assert cam_data["resolution"] == (20, 10)  # (W, H)

    assert "pooled_depth" in cam_data
    assert "pooled_confidence" in cam_data
    assert "pool_metadata" in cam_data
    assert "raw_frames" in cam_data

    assert len(cam_data["raw_frames"]) == 2
    assert cam_data["raw_frames"][0]["frame_index"] == 10
    assert cam_data["raw_frames"][1]["frame_index"] == 20

    # Check that raw frames contain depth maps
    assert "depth_map" in cam_data["raw_frames"][0]
    np.testing.assert_array_equal(cam_data["raw_frames"][0]["depth_map"], f1.depth_map)


def test_save_depth_skipped_when_no_path(mock_dependencies):
    """
    Test that save_depth_data is NOT called when save_depth_path is None.
    """
    _, _, _, mock_save_depth = mock_dependencies

    serial = "123456"
    serial_int = int(serial)
    results = {serial: {"pose": "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"}}

    f1 = MagicMock()
    f1.depth_map = np.ones((10, 10))
    f1.confidence_map = np.zeros((10, 10))

    vfs = [
        {
            "frame": f1,
            "ids": np.array([[1]]),
            "corners": np.zeros((1, 4, 2)),
            "score": 100,
            "frame_index": 10,
        }
    ]

    apply_depth_verify_refine_postprocess(
        results=results,
        verification_frames={serial_int: vfs},
        marker_geometry={1: np.zeros((4, 3))},
        camera_matrices={serial_int: np.eye(3)},
        verify_depth=True,
        refine_depth=False,
        use_confidence_weights=False,
        depth_confidence_threshold=50,
        save_depth_path=None,
    )

    mock_save_depth.assert_not_called()