import json
from pathlib import Path

import numpy as np
import pytest
import torch
from PIL import Image

from mcss.config import DatasetConfig
from mcss.data.manifest import ManifestFrame, SceneManifest, save_manifest
from mcss.data.manifest_dataset import ManifestSceneDataset, _find_manifests
from mcss.engine import build_dataset
from mcss.geometry import generate_rays


def _write_minimal_manifest(root: Path, scene_id: str) -> None:
    root.mkdir(parents=True)
    frames = tuple(
        ManifestFrame(
            frame_id=index,
            rgb=f"rgb/{index:06d}.png",
            intrinsics=[[2.0, 0.0, 0.5], [0.0, 2.0, 0.5], [0.0, 0.0, 1.0]],
            c2w=np.eye(4).tolist(),
        )
        for index in range(3)
    )
    save_manifest(
        SceneManifest(
            scene_id,
            (2, 2),
            ((-1.0, -1.0, -1.0), (1.0, 1.0, 1.0)),
            frames,
            root,
        ),
        root / "manifest.json",
    )


def test_manifest_dataset_loads_and_resizes_modalities(tmp_path: Path) -> None:
    scene = tmp_path / "scene"
    (scene / "rgb").mkdir(parents=True)
    (scene / "depth").mkdir()
    (scene / "normal").mkdir()
    for index in range(3):
        Image.new("RGB", (4, 4), (20 + index, 30, 40)).save(scene / "rgb" / f"{index:06d}.png")
        np.save(scene / "depth" / f"{index:06d}.npy", np.full((4, 4), 2.0, dtype=np.float32))
        normal = np.zeros((4, 4, 3), dtype=np.float32)
        normal[..., 2] = 1.0
        np.save(scene / "normal" / f"{index:06d}.npy", normal)
    frames = tuple(
        ManifestFrame(
            frame_id=index,
            rgb=f"rgb/{index:06d}.png",
            intrinsics=[[4.0, 0.0, 1.5], [0.0, 4.0, 1.5], [0.0, 0.0, 1.0]],
            c2w=np.eye(4).tolist(),
            depth=f"depth/{index:06d}.npy",
            normal=f"normal/{index:06d}.npy",
        )
        for index in range(3)
    )
    save_manifest(
        SceneManifest("scene", (4, 4), ((-1.0, -1.0, -1.0), (1.0, 1.0, 1.0)), frames, scene),
        scene / "manifest.json",
    )

    dataset = ManifestSceneDataset(scene, context_views=2, target_views=1, image_size=(2, 2))
    example = dataset[0]

    assert example.context_rgb.shape == (2, 3, 2, 2)
    assert example.target_depth.shape == (1, 1, 2, 2)
    assert example.target_normal.shape == (1, 3, 2, 2)
    assert example.target_point.shape == (1, 3, 2, 2)
    assert example.target_visibility.all()
    assert example.target_cameras.intrinsics[0, 0, 0] == 2.0
    assert example.target_cameras.intrinsics[0, 0, 2] == 0.75


def test_manifest_discovery_uses_direct_scene_fast_path(tmp_path: Path, monkeypatch) -> None:
    expected = []
    for name in ("scene_b", "scene_a"):
        scene = tmp_path / name
        scene.mkdir()
        manifest = scene / "manifest.json"
        manifest.write_text("{}", encoding="utf-8")
        expected.append(manifest)

    def fail_recursive_scan(_self: Path, _pattern: str):
        raise AssertionError("direct scene manifests should not trigger a recursive scan")

    monkeypatch.setattr(Path, "rglob", fail_recursive_scan)

    assert _find_manifests(tmp_path) == sorted(expected)


def test_manifest_dataset_filters_exact_scene_id_allowlist(tmp_path: Path) -> None:
    for scene_id in ("scene_c", "scene_a", "scene_b"):
        _write_minimal_manifest(tmp_path / scene_id, scene_id)

    dataset = ManifestSceneDataset(
        tmp_path,
        context_views=2,
        target_views=1,
        scene_ids=("scene_b", "scene_a"),
    )

    assert tuple(manifest.scene_id for manifest in dataset.manifests) == ("scene_b", "scene_a")
    assert len(dataset.entries) == 2
    with pytest.raises(ValueError, match="scene_missing"):
        ManifestSceneDataset(
            tmp_path,
            context_views=2,
            target_views=1,
            scene_ids=("scene_a", "scene_missing"),
        )


def test_build_dataset_passes_scene_id_allowlist(tmp_path: Path) -> None:
    for scene_id in ("scene_a", "scene_b"):
        _write_minimal_manifest(tmp_path / scene_id, scene_id)
    config = DatasetConfig(
        name="manifest",
        root=str(tmp_path),
        image_size=(2, 2),
        context_views=2,
        target_views=1,
        scene_ids=("scene_b",),
    )

    dataset = build_dataset(config, seed=3)

    assert isinstance(dataset, ManifestSceneDataset)
    assert tuple(manifest.scene_id for manifest in dataset.manifests) == ("scene_b",)


def _write_local_protocol_fixture(root: Path) -> Path:
    scene = root / "local_scene"
    (scene / "rgb").mkdir(parents=True)
    (scene / "depth").mkdir()
    (scene / "normal").mkdir()
    rotation = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float32)
    poses = []
    for translation in ([10.0, 0.0, 0.0], [10.0, 0.5, 0.0], [10.0, 1.0, 0.0]):
        pose = np.eye(4, dtype=np.float32)
        pose[:3, :3] = rotation
        pose[:3, 3] = translation
        poses.append(pose)
    frames = []
    for index, pose in enumerate(poses):
        Image.new("RGB", (2, 2), (20 + index, 30, 40)).save(scene / "rgb" / f"{index:06d}.png")
        depth_path = None
        normal_path = None
        if index == 2:
            np.save(scene / "depth" / f"{index:06d}.npy", np.full((2, 2), 2.0, np.float32))
            world_normal = np.zeros((2, 2, 3), dtype=np.float32)
            world_normal[..., 1] = 1.0
            np.save(scene / "normal" / f"{index:06d}.npy", world_normal)
            depth_path = f"depth/{index:06d}.npy"
            normal_path = f"normal/{index:06d}.npy"
        frames.append(
            ManifestFrame(
                frame_id=index,
                rgb=f"rgb/{index:06d}.png",
                intrinsics=[[2.0, 0.0, 0.5], [0.0, 2.0, 0.5], [0.0, 0.0, 1.0]],
                c2w=pose.tolist(),
                depth=depth_path,
                normal=normal_path,
            )
        )
    save_manifest(
        SceneManifest(
            "local_scene",
            (2, 2),
            ((100.0, 100.0, 100.0), (200.0, 200.0, 200.0)),
            tuple(frames),
            scene,
        ),
        scene / "manifest.json",
    )
    return scene


def test_context_local_protocol_ignores_manifest_and_target_geometry_for_support(
    tmp_path: Path,
) -> None:
    scene = _write_local_protocol_fixture(tmp_path)
    local_bounds = ((-2.0, -2.0, 0.1), (2.0, 2.0, 4.1))
    dataset = ManifestSceneDataset(
        scene,
        context_views=2,
        target_views=1,
        spatial_protocol="context_local_metric",
        local_bounds_m=local_bounds,
    )

    first = dataset[0]
    np.save(scene / "depth" / "000002.npy", np.full((2, 2), 200.0, np.float32))
    changed_target = dataset[0]

    torch.testing.assert_close(first.bounds, torch.tensor(local_bounds))
    torch.testing.assert_close(first.bounds, changed_target.bounds)
    torch.testing.assert_close(first.target_support, changed_target.target_support)
    assert first.target_visibility.all()
    assert not changed_target.target_visibility.any()
    assert not changed_target.target_depth.any()
    assert not changed_target.target_normal.any()
    assert not changed_target.target_point.any()
    assert not changed_target.target_rgb.any()
    assert first.scene_id == "local_scene"


def test_context_local_protocol_rebases_cameras_points_and_world_normals(tmp_path: Path) -> None:
    scene = _write_local_protocol_fixture(tmp_path)
    dataset = ManifestSceneDataset(
        scene,
        context_views=2,
        target_views=1,
        spatial_protocol="context_local_metric",
        local_bounds_m=((-2.0, -2.0, 0.1), (2.0, 2.0, 4.1)),
    )

    example = dataset[0]
    origins, directions = generate_rays(example.target_cameras)
    expected_points = origins + directions * example.target_depth[:, 0].unsqueeze(-1)

    torch.testing.assert_close(example.context_cameras.c2w[0], torch.eye(4), atol=1e-6, rtol=1e-6)
    torch.testing.assert_close(example.target_cameras.c2w[0, :3, 3], torch.tensor([1.0, 0.0, 0.0]))
    torch.testing.assert_close(
        example.target_point, expected_points.permute(0, 3, 1, 2), atol=1e-6, rtol=1e-6
    )
    expected_normal = (
        torch.tensor([1.0, 0.0, 0.0]).view(1, 3, 1, 1).expand_as(example.target_normal)
    )
    torch.testing.assert_close(example.target_normal, expected_normal, atol=1e-6, rtol=1e-6)
    assert example.target_support.shape == (1, 1, 2, 2)


def test_context_local_protocol_preserves_explicit_metric_point_magnitude(tmp_path: Path) -> None:
    scene = _write_local_protocol_fixture(tmp_path)
    (scene / "point").mkdir()
    world_points = np.empty((2, 2, 3), dtype=np.float32)
    world_points[...] = [10.0, 3.0, 2.0]
    np.save(scene / "point" / "000002.npy", world_points)
    manifest_path = scene / "manifest.json"
    manifest_data = json.loads(manifest_path.read_text(encoding="utf-8"))
    manifest_data["frames"][2]["point"] = "point/000002.npy"
    manifest_path.write_text(json.dumps(manifest_data), encoding="utf-8")
    dataset = ManifestSceneDataset(
        scene,
        context_views=2,
        target_views=1,
        spatial_protocol="context_local_metric",
        local_bounds_m=((-4.0, -2.0, 0.1), (4.0, 2.0, 4.1)),
    )

    example = dataset[0]

    expected = torch.tensor([3.0, 0.0, 2.0]).view(1, 3, 1, 1).expand_as(example.target_point)
    torch.testing.assert_close(example.target_point, expected, atol=1e-6, rtol=1e-6)