"""Tests for the unwarp/valid-mask ("UV domain") map the export builds.

Two layers:

* the module itself (``orisys_sensor_design.uv_map``) — needs only numpy, so it runs with
  the interpreter that builds the SDK-side artifacts;
* the export flow in ``cli.cmd_export`` — needs PyYAML for the profile loader, so those
  tests skip when it is absent rather than failing for the wrong reason.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest

from orisys_sensor_design import uv_map as uv

REPO_ROOT = Path(__file__).resolve().parents[1]
REVISION = REPO_ROOT / "sensors" / "finger" / "r01"
COMPUTE_PROFILE = REVISION / "profiles" / "compute.yaml"

#: A 1280x720 export like sensors/finger/r01/.../compute/camera_intrinsics.npz.
NATIVE_K = np.array(
    [[415.62082888, 0.0, 640.0], [0.0, 415.62082888, 360.0], [0.0, 0.0, 1.0]]
)


# ----------------------------------------------------------------------
# Derived source scale
# ----------------------------------------------------------------------
def test_derive_source_scale_follows_the_output_requirement() -> None:
    # max(280/1280, 300/720) = 0.4166... -> 0.417; the floor that avoids upsampling.
    assert uv.derive_source_scale(1280, 720, 280, 300) == 0.417


def test_derive_source_scale_is_camera_independent_for_one_output_grid() -> None:
    # Both cameras land on the same 533x300 source grid because the rule follows the
    # output requirement, not the native frame.
    assert uv.derive_source_scale(1920, 1080, 280, 300) == 0.278
    scale = uv.derive_source_scale(1920, 1080, 280, 300)
    assert int(1920 * scale) >= 280 and int(1080 * scale) >= 300


def test_derive_source_scale_covers_the_output_one_to_one() -> None:
    assert uv.derive_source_scale(1280, 720, 640, 360) == 0.5
    # A downscale to half the frame still needs half the scale, not 1.0.
    assert uv.derive_source_scale(1280, 720, 280, 300) < 0.5


def test_derive_source_scale_never_exceeds_the_native_frame() -> None:
    # An output grid larger than the camera cannot be covered without upsampling, so the
    # scale is clamped rather than raised above 1.0.
    assert uv.derive_source_scale(320, 240, 640, 480) == 1.0


@pytest.mark.parametrize(
    "native_width, native_height",
    [(0, 720), (1280, 0), (-1, 720)],
)
def test_derive_source_scale_rejects_a_non_positive_camera(
    native_width: int, native_height: int
) -> None:
    with pytest.raises(uv.UvMapError, match="non-positive"):
        uv.derive_source_scale(native_width, native_height, 280, 300)


# ----------------------------------------------------------------------
# Parameter resolution: never guess
# ----------------------------------------------------------------------
GOOD_BLOCK = {
    "out_width": 280,
    "out_height": 300,
    "flip_v": True,
    "source_scale": 0.4,
    "edge_band": 0.12,
    "edge_expand": 1.0,
}


def test_load_uv_map_params_reads_and_types_the_block() -> None:
    block = dict(GOOD_BLOCK, out_width="280", source_scale="0.4")
    params = uv.load_uv_map_params({"uv_map": block})
    assert params == {
        "out_width": 280,
        "out_height": 300,
        "flip_v": True,
        "source_scale": 0.4,
        "edge_band": 0.12,
        "edge_expand": 1.0,
    }
    assert isinstance(params["out_width"], int)
    assert isinstance(params["flip_v"], bool)


def test_load_uv_map_params_errors_without_a_uv_map_block() -> None:
    # The whole point: a missing block must fail, not fall back to a default domain.
    with pytest.raises(uv.UvMapError, match="no uv_map: block"):
        uv.load_uv_map_params({"profile": "compute", "artifacts": ["mesh.npz"]})


def test_load_uv_map_params_errors_on_an_empty_block() -> None:
    with pytest.raises(uv.UvMapError, match="no uv_map: block"):
        uv.load_uv_map_params({"uv_map": {}})


def test_load_uv_map_params_errors_on_a_missing_key() -> None:
    block = dict(GOOD_BLOCK)
    del block["source_scale"]
    with pytest.raises(uv.UvMapError, match="missing \\['source_scale'\\]"):
        uv.load_uv_map_params({"uv_map": block})


def test_load_uv_map_params_errors_on_an_unknown_key() -> None:
    block = dict(GOOD_BLOCK, out_depth=4)
    with pytest.raises(uv.UvMapError, match="does not understand"):
        uv.load_uv_map_params({"uv_map": block})


def test_quoted_false_does_not_read_as_true() -> None:
    params = uv.load_uv_map_params({"uv_map": dict(GOOD_BLOCK, flip_v="false")})
    assert params["flip_v"] is False


def test_compute_profile_pins_the_scale_the_map_is_built_at() -> None:
    pytest.importorskip("yaml", reason="reading the profile needs PyYAML")
    from orisys.unroll.geometry import scaled_map_cache_path

    from orisys_sensor_design.manifest import load_yaml

    profile = load_yaml(COMPUTE_PROFILE)
    params = uv.load_uv_map_params(profile)
    # Whatever the profile pins is what the build uses (the value is a deliberate choice,
    # so this asserts the wiring, not one particular number).
    assert params["source_scale"] == float(profile["uv_map"]["source_scale"])
    assert (params["out_width"], params["out_height"]) == (280, 300)
    # ... and the artifact the export must produce is named for that same scale, so a
    # changed scale cannot leave a stale, differently-named map looking complete.
    expected = Path(scaled_map_cache_path("unroll_map_uv.npz", params["source_scale"])).name
    assert expected in profile["artifacts"]


# ----------------------------------------------------------------------
# build_uv_map: the SDK call, with a stubbed SDK
# ----------------------------------------------------------------------
@pytest.fixture()
def export_dir(tmp_path: Path) -> Path:
    """An export payload directory with only the calibration the map step reads."""
    directory = tmp_path / "compute"
    directory.mkdir()
    np.savez(
        directory / "camera_intrinsics.npz",
        K=NATIVE_K,
        width=np.array(1280),
        height=np.array(720),
    )
    (directory / "mesh.npz").write_bytes(b"mesh")
    (directory / "world_to_cam.npy").write_bytes(b"pose")
    return directory


def _stub_sdk(monkeypatch: pytest.MonkeyPatch, calls: dict, mask_mode: str = "valid"):
    """Replace the two SDK entry points the map step calls, recording the arguments."""
    import orisys.unroll.geometry as geometry
    import orisys.unroll.precompute as precompute

    def fake_scale_camera_geometry(K, width, height, scale):
        calls["scale_args"] = (width, height, scale)
        return np.asarray(K, dtype=np.float64) * scale, int(round(width * scale)), int(
            round(height * scale)
        )

    def fake_build_or_load_domain(mesh, intrinsics, pose, map_path, **kwargs):
        calls["build_kwargs"] = kwargs
        calls["target"] = map_path
        calls["target_existed"] = Path(map_path).is_file()
        width = int(kwargs["out_width"])
        height = int(kwargs["out_height"])
        arrays = {
            "map_x": np.zeros((height, width), np.float32),
            "map_y": np.zeros((height, width), np.float32),
        }
        if mask_mode == "valid":
            mask = np.zeros((height, width), bool)
            mask[height // 2 :, :] = True
            arrays["valid_mask"] = mask
            arrays["world_xyz"] = np.zeros((height, width, 3), np.float32)
        elif mask_mode == "empty":
            arrays["valid_mask"] = np.zeros((height, width), bool)
            arrays["world_xyz"] = np.zeros((height, width, 3), np.float32)
        elif mask_mode == "no_surface":
            arrays["valid_mask"] = np.ones((height, width), bool)
        elif mask_mode == "wrong_shape":
            arrays["valid_mask"] = np.ones((4, 4), bool)
            arrays["world_xyz"] = np.zeros((4, 4, 3), np.float32)
        np.savez(map_path, **arrays)

    monkeypatch.setattr(geometry, "scale_camera_geometry", fake_scale_camera_geometry)
    monkeypatch.setattr(precompute, "build_or_load_domain", fake_build_or_load_domain)


def test_build_uv_map_scales_source_K_to_the_profile_scale(
    export_dir: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    calls: dict = {}
    _stub_sdk(monkeypatch, calls)
    params = uv.load_uv_map_params({"uv_map": GOOD_BLOCK})

    target = uv.build_uv_map(export_dir, params, export_dir)

    assert target == export_dir / "unroll_map_uv_s0.4.npz"
    assert target.is_file()
    kwargs = calls["build_kwargs"]
    # 1280 x 0.4 — the scaled K, not the native one (that measured 3.5 % instead of 40.8 %).
    assert kwargs["source_width"] == 512
    assert kwargs["source_height"] == 288
    assert kwargs["source_scale"] == 0.4
    assert kwargs["source_K"][0, 0] == pytest.approx(NATIVE_K[0, 0] * 0.4)
    assert (kwargs["out_width"], kwargs["out_height"]) == (280, 300)
    assert kwargs["include_surface"] is True


def test_build_uv_map_takes_the_source_frame_from_this_export(
    export_dir: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    calls: dict = {}
    _stub_sdk(monkeypatch, calls)
    params = uv.load_uv_map_params({"uv_map": GOOD_BLOCK})

    uv.build_uv_map(export_dir, params, export_dir)

    # The source frame is derived from *this* export's camera_intrinsics.npz, never
    # inherited from another camera's map (that measured 0.8 %).
    assert calls["scale_args"] == (1280, 720, 0.4)


def test_build_uv_map_derives_the_scale_when_none_is_pinned(
    export_dir: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    calls: dict = {}
    _stub_sdk(monkeypatch, calls)
    params = dict(GOOD_BLOCK)
    del params["source_scale"]

    target = uv.build_uv_map(export_dir, params, export_dir)

    assert target.name == "unroll_map_uv_s0.417.npz"
    assert calls["build_kwargs"]["source_scale"] == 0.417


def test_force_deletes_the_target_so_the_sdk_must_rebuild(
    export_dir: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    calls: dict = {}
    _stub_sdk(monkeypatch, calls)
    params = uv.load_uv_map_params({"uv_map": GOOD_BLOCK})
    stale = export_dir / "unroll_map_uv_s0.4.npz"
    stale.write_bytes(b"stale map")

    uv.build_uv_map(export_dir, params, export_dir)
    assert calls["target_existed"] is True  # without --force the SDK may load it

    uv.build_uv_map(export_dir, params, export_dir, force=True)
    assert calls["target_existed"] is False


def test_build_uv_map_refuses_a_payload_the_export_did_not_write(tmp_path: Path) -> None:
    empty = tmp_path / "empty"
    empty.mkdir()
    with pytest.raises(uv.UvMapError, match="missing"):
        uv.build_uv_map(empty, dict(GOOD_BLOCK), empty)


@pytest.mark.parametrize(
    "mask_mode, message",
    [
        ("no_surface", "world_xyz"),
        ("empty", "no set pixel"),
        ("wrong_shape", "does not match"),
    ],
)
def test_build_uv_map_rejects_an_incomplete_map(
    export_dir: Path, monkeypatch: pytest.MonkeyPatch, mask_mode: str, message: str
) -> None:
    calls: dict = {}
    _stub_sdk(monkeypatch, calls, mask_mode=mask_mode)
    params = uv.load_uv_map_params({"uv_map": GOOD_BLOCK})

    with pytest.raises(uv.UvMapError, match=message):
        uv.build_uv_map(export_dir, params, export_dir)


def test_a_map_without_a_valid_mask_is_fatal(
    export_dir: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    # mask_mode="valid" always writes one; drop it by hand to model an old/incomplete cache.
    calls: dict = {}
    _stub_sdk(monkeypatch, calls)
    params = uv.load_uv_map_params({"uv_map": GOOD_BLOCK})

    import orisys.unroll.precompute as precompute

    def build_without_mask(mesh, intrinsics, pose, map_path, **kwargs):
        calls["build_kwargs"] = kwargs
        height, width = int(kwargs["out_height"]), int(kwargs["out_width"])
        np.savez(map_path, map_x=np.zeros((height, width), np.float32))

    monkeypatch.setattr(precompute, "build_or_load_domain", build_without_mask)
    with pytest.raises(uv.UvMapError, match="valid_mask"):
        uv.build_uv_map(export_dir, params, export_dir)


# ----------------------------------------------------------------------
# Export flow: cmd_export builds the map between Blender and the artifact check
# ----------------------------------------------------------------------
def _cli():
    pytest.importorskip("yaml", reason="the CLI's profile loader needs PyYAML")
    from orisys_sensor_design import cli

    return cli


def _prepare_export(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, profile: dict | None = None):
    """Point cmd_export at a temp output dir with the Blender artifacts already written."""
    cli = _cli()
    output_dir = tmp_path / "compute"
    output_dir.mkdir()
    for name in ("mesh.npz", "housing.npz", "camera_intrinsics.npz", "world_to_cam.npy"):
        (output_dir / name).write_bytes(b"artifact")
    monkeypatch.setattr(cli, "export_output_dir", lambda *a, **k: output_dir)
    monkeypatch.setattr(cli, "export_revision_dir", lambda *a, **k: tmp_path / "finger" / "r01")
    blender_calls: list = []
    monkeypatch.setattr(cli, "_run_blender_script", lambda *a, **k: blender_calls.append(a) or 0)
    if profile is not None:
        monkeypatch.setattr(cli, "load_profile", lambda paths, name: profile)
    return cli, output_dir, blender_calls


def _record_builder(monkeypatch: pytest.MonkeyPatch, recorded: dict, *, fail: bool = False):
    from orisys.unroll.geometry import scaled_map_cache_path

    def fake_build(export_dir, params, output_dir, *, force=False):
        if fail:
            raise uv.UvMapError("the island did not land in the output domain")
        recorded.update(
            export_dir=Path(export_dir),
            params=params,
            output_dir=Path(output_dir),
            force=force,
        )
        # Name it the way the real builder does, so the artifact check sees what the
        # profile promises regardless of the scale the profile pins.
        target = Path(
            scaled_map_cache_path(
                str(Path(output_dir) / "unroll_map_uv.npz"), float(params["source_scale"])
            )
        )
        target.write_bytes(b"map")
        return target

    monkeypatch.setattr("orisys_sensor_design.uv_map.build_uv_map", fake_build)


def test_export_builds_the_map_in_the_same_run(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
) -> None:
    cli, output_dir, blender_calls = _prepare_export(tmp_path, monkeypatch)
    recorded: dict = {}
    _record_builder(monkeypatch, recorded)
    from orisys_sensor_design.manifest import load_yaml

    expected = uv.load_uv_map_params(load_yaml(COMPUTE_PROFILE))

    code = cli.cmd_export(REVISION, "compute")

    assert code == 0
    assert blender_calls, "the Blender export must still run first"
    assert recorded["export_dir"] == output_dir
    assert recorded["output_dir"] == output_dir
    # Exactly the profile's uv_map: parameters - including the pinned source_scale.
    assert recorded["params"] == expected
    out = capsys.readouterr().out
    assert "OK: all profile artifacts present" in out


def test_export_returns_a_distinct_code_when_the_map_fails(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
) -> None:
    cli, _output_dir, _blender = _prepare_export(tmp_path, monkeypatch)
    _record_builder(monkeypatch, {}, fail=True)

    code = cli.cmd_export(REVISION, "compute")

    captured = capsys.readouterr()
    assert code == cli.UV_MAP_EXIT_CODE == 4
    assert "unwarp/valid-mask" in captured.err
    # The bundle must not be reported as complete: the map is what stops it rendering grey.
    assert "OK: all profile artifacts present" not in captured.out


def test_export_reports_a_profile_that_requires_a_map_it_cannot_build(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
) -> None:
    profile = {"profile": "compute", "artifacts": ["mesh.npz", "unroll_map_uv_s0.4.npz"]}
    cli, _output_dir, _blender = _prepare_export(tmp_path, monkeypatch, profile=profile)
    recorded: dict = {}
    _record_builder(monkeypatch, recorded)

    code = cli.cmd_export(REVISION, "compute")

    captured = capsys.readouterr()
    assert code == cli.UV_MAP_EXIT_CODE
    assert "no uv_map: block" in captured.err
    assert recorded == {}, "the builder must not be called without parameters"


def test_export_skips_the_map_step_for_a_profile_that_declares_none(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture
) -> None:
    # Profiles whose runtime bundle is still SDK-owned keep their old behaviour.
    profile = {"profile": "compute", "artifacts": ["mesh.npz"]}
    cli, _output_dir, _blender = _prepare_export(tmp_path, monkeypatch, profile=profile)
    recorded: dict = {}
    _record_builder(monkeypatch, recorded)

    code = cli.cmd_export(REVISION, "compute")

    captured = capsys.readouterr()
    assert code == 0
    assert recorded == {}
    assert "skipping the unwarp/valid-mask build" in captured.out
    assert "OK: all profile artifacts present" in captured.out