"""Tests for metadata-only full train/dev episode compilation."""

from __future__ import annotations

import importlib.util
import json
import sys
from pathlib import Path

import numpy as np
import pytest

SPEC = importlib.util.spec_from_file_location(
    "compile_dynamic_training_episodes",
    Path(__file__).resolve().parents[1] / "scripts/compile_dynamic_training_episodes.py",
)
assert SPEC is not None and SPEC.loader is not None
compiler = importlib.util.module_from_spec(SPEC)
sys.modules[SPEC.name] = compiler
SPEC.loader.exec_module(compiler)


def _write_scene(
    root: Path,
    split: str,
    scene: str,
    camera: str,
    count: int = 16,
    invalid_camera_frames: tuple[int, ...] = (),
) -> None:
    scene_root = root / split / scene
    (scene_root / "rgb").mkdir(parents=True)
    (scene_root / "depth").mkdir()
    frames = []
    for frame_id in range(count):
        (scene_root / "rgb" / f"{frame_id:06d}.png").write_bytes(b"not an image")
        (scene_root / "depth" / f"{frame_id:06d}.npy").write_bytes(b"not an array")
        c2w = np.eye(4).tolist()
        if frame_id in invalid_camera_frames:
            if frame_id % 2 == 0:
                c2w[0][0] = 1.01
            else:
                c2w[0][0] = -1.0
        frames.append(
            {
                "frame_id": frame_id,
                "rgb": f"rgb/{frame_id:06d}.png",
                "depth": f"depth/{frame_id:06d}.npy",
                "normal": f"normal/{frame_id:06d}.npy",
                "intrinsics": [[2.0, 0.0, 1.5], [0.0, 2.0, 1.5], [0.0, 0.0, 1.0]],
                "c2w": c2w,
            }
        )
    payload = {
        "schema_version": 1,
        "scene_id": scene,
        "coordinate_convention": "opencv_c2w",
        "image_size": [4, 6],
        "bounds": [[-1.0, -1.0, -1.0], [1.0, 1.0, 1.0]],
        "frames": frames[::-1],
    }
    (scene_root / "manifest.json").write_text(json.dumps(payload), encoding="utf-8")


def _write_partitions(path: Path, rows: list[tuple[str, str, str, int]]) -> None:
    lines = [
        "scene_name,official_split,protocol_partition,selected_camera,frame_count",
        *(f"{scene},{split},{split},{camera},{count}" for scene, split, camera, count in rows),
    ]
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def test_compiles_fixed_segments_and_does_not_decode_rgb_or_depth(tmp_path: Path) -> None:
    prepared = tmp_path / "prepared"
    _write_scene(prepared, "train", "scene_train", "cam_01")
    _write_scene(prepared, "val", "scene_val", "cam_00")
    partitions = tmp_path / "partitions.csv"
    _write_partitions(
        partitions,
        [("scene_train", "train", "cam_01", 16), ("scene_val", "val", "cam_00", 16)],
    )
    inventory = compiler.compile_training_episodes(
        prepared_root=prepared,
        partitions_path=partitions,
        output_root=tmp_path / "output",
        expected_counts=None,
        train_depth_health_scan=False,
    )
    assert inventory["splits"]["train"]["eligible_episode_count"] == 1
    assert inventory["splits"]["dev"]["eligible_episode_count"] == 1
    train_index = json.loads(
        (tmp_path / "output/episodes/episode_index_train.json").read_text(encoding="utf-8")
    )
    train_online = json.loads(
        Path(train_index["episodes"][0]["online_manifest"]).read_text(encoding="utf-8")
    )
    train_query = json.loads(
        Path(train_index["episodes"][0]["query_manifest"]).read_text(encoding="utf-8")
    )
    assert train_online["episode_id"] == "train--scene_train--cam_01"
    assert [frame["frame_id"] for frame in train_online["warmup"]] == [0, 1, 2, 3]
    assert [frame["frame_id"] for frame in train_online["stream"]] == list(range(4, 12))
    assert [frame["frame_id"] for frame in train_query["query"]] == [12, 13, 14, 15]
    assert all("depth" not in frame for frame in train_online["warmup"] + train_online["stream"])
    assert all("normal" not in frame for frame in train_query["query"])


def test_short_scene_fails_by_default_and_is_explicit_when_allowed(tmp_path: Path) -> None:
    prepared = tmp_path / "prepared"
    _write_scene(prepared, "train", "scene_short", "cam_00", count=9)
    (prepared / "val").mkdir(parents=True)
    partitions = tmp_path / "partitions.csv"
    _write_partitions(partitions, [("scene_short", "train", "cam_00", 9)])
    with pytest.raises(ValueError, match="ineligible scenes"):
        compiler.compile_training_episodes(
            prepared_root=prepared,
            partitions_path=partitions,
            output_root=tmp_path / "strict",
            expected_counts=None,
            train_depth_health_scan=False,
        )
    inventory = compiler.compile_training_episodes(
        prepared_root=prepared,
        partitions_path=partitions,
        output_root=tmp_path / "allowed",
        expected_counts=None,
        allow_ineligible=True,
        train_depth_health_scan=False,
    )
    assert inventory["ineligible_scene_ids"] == ["scene_short"]
    assert inventory["splits"]["train"]["eligible_episode_count"] == 0


def test_missing_selected_reference_fails(tmp_path: Path) -> None:
    prepared = tmp_path / "prepared"
    _write_scene(prepared, "train", "scene_train", "cam_00")
    (prepared / "val").mkdir(parents=True)
    (prepared / "train/scene_train/depth/000015.npy").unlink()
    partitions = tmp_path / "partitions.csv"
    _write_partitions(partitions, [("scene_train", "train", "cam_00", 16)])
    with pytest.raises(FileNotFoundError, match="missing .*depth reference"):
        compiler.compile_training_episodes(
            prepared_root=prepared,
            partitions_path=partitions,
            output_root=tmp_path / "output",
            expected_counts=None,
            train_depth_health_scan=False,
        )


def test_train_val_scene_intersection_fails(tmp_path: Path) -> None:
    prepared = tmp_path / "prepared"
    _write_scene(prepared, "train", "same_scene", "cam_00")
    _write_scene(prepared, "val", "same_scene", "cam_00")
    partitions = tmp_path / "partitions.csv"
    _write_partitions(
        partitions,
        [("same_scene", "train", "cam_00", 16), ("same_scene", "val", "cam_00", 16)],
    )
    with pytest.raises(ValueError, match="duplicate train/val partition row"):
        compiler.compile_training_episodes(
            prepared_root=prepared,
            partitions_path=partitions,
            output_root=tmp_path / "output",
            expected_counts=None,
            train_depth_health_scan=False,
        )


def test_compiles_multiple_chronological_windows_for_one_train_scene(tmp_path: Path) -> None:
    prepared = tmp_path / "prepared"
    _write_scene(prepared, "train", "scene_train", "cam_00", count=32)
    _write_scene(prepared, "val", "scene_val", "cam_00", count=16)
    partitions = tmp_path / "partitions.csv"
    _write_partitions(
        partitions,
        [
            ("scene_train", "train", "cam_00", 32),
            ("scene_val", "val", "cam_00", 16),
        ],
    )
    inventory = compiler.compile_training_episodes(
        prepared_root=prepared,
        partitions_path=partitions,
        output_root=tmp_path / "output",
        expected_counts=None,
        train_depth_health_scan=False,
        train_window_stride=16,
        max_train_windows_per_scene=None,
    )
    assert inventory["splits"]["train"]["eligible_episode_count"] == 2
    train_index = json.loads(
        (tmp_path / "output/episodes/episode_index_train.json").read_text(encoding="utf-8")
    )
    assert len(train_index["episodes"]) == 2
    assert len({entry["episode_id"] for entry in train_index["episodes"]}) == 2
    windows = []
    for entry in train_index["episodes"]:
        online = json.loads(Path(entry["online_manifest"]).read_text(encoding="utf-8"))
        query = json.loads(Path(entry["query_manifest"]).read_text(encoding="utf-8"))
        observed = [frame["frame_id"] for frame in online["warmup"] + online["stream"]]
        queries = [frame["frame_id"] for frame in query["query"]]
        assert max(observed) < min(queries)
        assert set(observed).isdisjoint(queries)
        windows.append(observed + queries)
    assert windows == [list(range(16)), list(range(16, 32))]
    dev_index = json.loads(
        (tmp_path / "output/episodes/episode_index_dev.json").read_text(encoding="utf-8")
    )
    assert len(dev_index["episodes"]) == 1


def test_rigid_camera_filter_excludes_metadata_invalid_frames_and_records_hashes(
    tmp_path: Path,
) -> None:
    prepared = tmp_path / "prepared"
    _write_scene(
        prepared,
        "train",
        "scene_train",
        "cam_00",
        count=18,
        invalid_camera_frames=(0, 1),
    )
    _write_scene(prepared, "val", "scene_val", "cam_00", count=16)
    partitions = tmp_path / "partitions.csv"
    _write_partitions(
        partitions,
        [
            ("scene_train", "train", "cam_00", 18),
            ("scene_val", "val", "cam_00", 16),
        ],
    )

    inventory = compiler.compile_training_episodes(
        prepared_root=prepared,
        partitions_path=partitions,
        output_root=tmp_path / "output",
        expected_counts=None,
        train_depth_health_scan=False,
        rigid_cameras_only=True,
    )

    manifest_path = prepared / "train/scene_train/manifest.json"
    assert inventory["rigid_cameras_only"] is True
    assert inventory["rigid_camera_tolerance"] == pytest.approx(1e-3)
    assert inventory["rigid_camera_exclusions"] == [
        {
            "scene_id": "scene_train",
            "source_split": "train",
            "split_id": "train",
            "manifest": manifest_path.resolve().as_posix(),
            "source_manifest_sha256": compiler.sha256(manifest_path),
            "original_frame_count": 18,
            "excluded_frame_ids": [0, 1],
            "excluded_frames": [
                {"frame_id": 0, "reason": "rotation_orthogonality"},
                {"frame_id": 1, "reason": "rotation_determinant"},
            ],
        }
    ]
    train_index = json.loads(
        (tmp_path / "output/episodes/episode_index_train.json").read_text(encoding="utf-8")
    )
    online = json.loads(Path(train_index["episodes"][0]["online_manifest"]).read_text())
    query = json.loads(Path(train_index["episodes"][0]["query_manifest"]).read_text())
    assert [frame["frame_id"] for frame in online["warmup"] + online["stream"]] == list(
        range(2, 14)
    )
    assert [frame["frame_id"] for frame in query["query"]] == list(range(14, 18))
    assert all(frame["c2w"] == np.eye(4).tolist() for frame in online["warmup"])

    scene_inventory = json.loads(
        (tmp_path / "output/inventory/scene_inventory_train.json").read_text(encoding="utf-8")
    )["scenes"][0]
    assert scene_inventory["original_frame_count"] == 18
    assert scene_inventory["frame_count"] == 16
    assert scene_inventory["excluded_camera_frames"] == [
        {"frame_id": 0, "reason": "rotation_orthogonality"},
        {"frame_id": 1, "reason": "rotation_determinant"},
    ]


def test_rigid_camera_filter_does_not_silently_drop_short_scene(tmp_path: Path) -> None:
    prepared = tmp_path / "prepared"
    _write_scene(
        prepared,
        "train",
        "scene_train",
        "cam_00",
        count=16,
        invalid_camera_frames=(0,),
    )
    _write_scene(prepared, "val", "scene_val", "cam_00", count=16)
    partitions = tmp_path / "partitions.csv"
    _write_partitions(
        partitions,
        [
            ("scene_train", "train", "cam_00", 16),
            ("scene_val", "val", "cam_00", 16),
        ],
    )

    with pytest.raises(ValueError, match="ineligible scenes.*scene_train"):
        compiler.compile_training_episodes(
            prepared_root=prepared,
            partitions_path=partitions,
            output_root=tmp_path / "output",
            expected_counts=None,
            train_depth_health_scan=False,
            rigid_cameras_only=True,
        )


def test_rigid_camera_filter_checks_partition_count_before_filtering(tmp_path: Path) -> None:
    prepared = tmp_path / "prepared"
    _write_scene(
        prepared,
        "train",
        "scene_train",
        "cam_00",
        count=16,
        invalid_camera_frames=(0,),
    )
    (prepared / "val").mkdir(parents=True)
    partitions = tmp_path / "partitions.csv"
    _write_partitions(partitions, [("scene_train", "train", "cam_00", 15)])

    with pytest.raises(ValueError, match="partition frame_count mismatch"):
        compiler.compile_training_episodes(
            prepared_root=prepared,
            partitions_path=partitions,
            output_root=tmp_path / "output",
            expected_counts=None,
            train_depth_health_scan=False,
            rigid_cameras_only=True,
        )