"""Tests for orbit.loaders — pluggable dataset loading system."""

from __future__ import annotations

import json

import numpy as np
import pandas as pd
import pytest

from orbit.analyzer.dataset_loader import DatasetInfo


# ---------------------------------------------------------------------------
# BaseLoader contract
# ---------------------------------------------------------------------------


class TestBaseLoader:
    def test_sample_episode_indices_no_sampling(self):
        from orbit.loaders.base import sample_episode_indices

        assert sample_episode_indices(5, 10) == [0, 1, 2, 3, 4]

    def test_sample_episode_indices_with_sampling(self):
        from orbit.loaders.base import sample_episode_indices

        result = sample_episode_indices(200, 10)
        assert len(result) == 10
        assert result[0] == 0
        assert all(isinstance(i, int) for i in result)

    def test_normalize_dataframe_adds_missing_columns(self):
        from orbit.loaders.base import normalize_dataframe

        df = pd.DataFrame({"episode_index": [0], "action": [[1.0, 2.0]]})
        result = normalize_dataframe(df)
        assert "state" in result.columns
        assert "timestamp" in result.columns
        assert "frame_index" in result.columns

    def test_normalize_dataframe_renames_observation_state(self):
        from orbit.loaders.base import normalize_dataframe

        df = pd.DataFrame(
            {
                "episode_index": [0],
                "frame_index": [0],
                "action": [[1.0]],
                "observation.state": [[2.0]],
                "timestamp": [0.0],
            }
        )
        result = normalize_dataframe(df)
        assert "state" in result.columns
        assert "observation.state" not in result.columns

    def test_normalize_empty_dataframe(self):
        from orbit.loaders.base import normalize_dataframe

        result = normalize_dataframe(pd.DataFrame())
        assert list(result.columns) == [
            "episode_index",
            "frame_index",
            "action",
            "state",
            "timestamp",
        ]


# ---------------------------------------------------------------------------
# Registry / Auto-detection
# ---------------------------------------------------------------------------


class TestDetection:
    def test_detect_lerobot_local(self, tmp_path):
        from orbit.loaders import detect_loader

        meta = tmp_path / "meta"
        meta.mkdir()
        (meta / "info.json").write_text(json.dumps({"fps": 30}))

        loader = detect_loader(str(tmp_path))
        assert loader.__class__.__name__ == "LeRobotLocalLoader"

    def test_detect_hdf5(self, tmp_path):
        pytest.importorskip("h5py")
        import h5py

        from orbit.loaders import detect_loader

        p = tmp_path / "demo.hdf5"
        with h5py.File(str(p), "w") as f:
            grp = f.create_group("data")
            demo = grp.create_group("demo_0")
            demo.create_dataset("actions", data=np.zeros((10, 7)))
            demo.create_dataset("states", data=np.zeros((10, 14)))

        loader = detect_loader(str(p))
        assert loader.__class__.__name__ == "HDF5Loader"

    def test_detect_directory_npz(self, tmp_path):
        from orbit.loaders import detect_loader

        np.savez(
            tmp_path / "episode_000.npz",
            action=np.zeros((10, 7)),
            state=np.zeros((10, 14)),
        )
        loader = detect_loader(str(tmp_path))
        assert loader.__class__.__name__ == "DirectoryLoader"

    def test_detect_hf_repo_id(self):
        from orbit.loaders import detect_loader

        loader = detect_loader("lerobot/pusht")
        assert loader.__class__.__name__ == "LeRobotHubLoader"

    def test_explicit_format_override(self):
        from orbit.loaders import detect_loader

        loader = detect_loader("anything", format="hdf5")
        assert loader.__class__.__name__ == "HDF5Loader"

    def test_unknown_format_raises(self):
        from orbit.loaders import detect_loader

        with pytest.raises(ValueError, match="Unknown format"):
            detect_loader("anything", format="unknown_format")

    def test_nonexistent_path_raises(self):
        from orbit.loaders import detect_loader

        with pytest.raises(ValueError):
            detect_loader("/nonexistent/path/to/dataset")


# ---------------------------------------------------------------------------
# LeRobot Local Loader
# ---------------------------------------------------------------------------


class TestLeRobotLocalLoader:
    @pytest.fixture()
    def lerobot_dir(self, tmp_path):
        """Create a minimal LeRobot v3 dataset directory."""
        meta = tmp_path / "meta"
        meta.mkdir()
        (meta / "info.json").write_text(
            json.dumps(
                {
                    "total_episodes": 3,
                    "total_frames": 150,
                    "fps": 50.0,
                    "robot_type": "so100",
                    "features": {
                        "action": {"dtype": "float32", "shape": [6]},
                        "observation.state": {"dtype": "float32", "shape": [6]},
                    },
                }
            )
        )

        # Create per-episode parquet files
        data_dir = tmp_path / "data" / "chunk-000"
        data_dir.mkdir(parents=True)

        for ep_idx in range(3):
            n_frames = 50
            df = pd.DataFrame(
                {
                    "episode_index": [ep_idx] * n_frames,
                    "frame_index": list(range(n_frames)),
                    "action": [np.random.randn(6).tolist() for _ in range(n_frames)],
                    "observation.state": [
                        np.random.randn(6).tolist() for _ in range(n_frames)
                    ],
                    "timestamp": [i / 50.0 for i in range(n_frames)],
                }
            )
            df.to_parquet(data_dir / f"episode_{ep_idx:06d}.parquet", index=False)

        return tmp_path

    def test_can_load(self, lerobot_dir):
        from orbit.loaders.lerobot_local import LeRobotLocalLoader

        loader = LeRobotLocalLoader()
        assert loader.can_load(str(lerobot_dir)) is True
        assert loader.can_load("/nonexistent") is False

    def test_load_info(self, lerobot_dir):
        from orbit.loaders.lerobot_local import LeRobotLocalLoader

        loader = LeRobotLocalLoader()
        info = loader.load_info(str(lerobot_dir))
        assert info.num_episodes == 3
        assert info.fps == 50.0
        assert info.robot_type == "so100"
        assert "action" in info.shapes

    def test_load_episodes(self, lerobot_dir):
        from orbit.loaders.lerobot_local import LeRobotLocalLoader

        loader = LeRobotLocalLoader()
        df = loader.load_episodes(str(lerobot_dir))
        assert not df.empty
        assert "episode_index" in df.columns
        assert "action" in df.columns
        assert "state" in df.columns
        assert df["episode_index"].nunique() == 3

    def test_load_specific_episodes(self, lerobot_dir):
        from orbit.loaders.lerobot_local import LeRobotLocalLoader

        loader = LeRobotLocalLoader()
        df = loader.load_episodes(str(lerobot_dir), episode_indices=[0, 2])
        assert df["episode_index"].nunique() == 2


# ---------------------------------------------------------------------------
# HDF5 Loader
# ---------------------------------------------------------------------------


class TestHDF5Loader:
    @pytest.fixture()
    def hdf5_file(self, tmp_path):
        """Create a minimal RoboMimic HDF5 file."""
        h5py = pytest.importorskip("h5py")

        p = tmp_path / "demo.hdf5"
        with h5py.File(str(p), "w") as f:
            data_grp = f.create_group("data")
            for i in range(5):
                demo = data_grp.create_group(f"demo_{i}")
                n = 20 + i * 5
                demo.create_dataset("actions", data=np.random.randn(n, 7).astype(np.float32))
                demo.create_dataset("states", data=np.random.randn(n, 14).astype(np.float32))
        return str(p)

    def test_can_load(self, hdf5_file):
        from orbit.loaders.hdf5 import HDF5Loader

        loader = HDF5Loader()
        assert loader.can_load(hdf5_file) is True
        assert loader.can_load("/nonexistent.hdf5") is False

    def test_load_info(self, hdf5_file):
        from orbit.loaders.hdf5 import HDF5Loader

        loader = HDF5Loader()
        info = loader.load_info(hdf5_file)
        assert info.num_episodes == 5
        assert info.num_frames > 0
        assert info.fps == 20.0
        assert info.shapes["action"] == [7]
        assert info.shapes["state"] == [14]

    def test_load_episodes(self, hdf5_file):
        from orbit.loaders.hdf5 import HDF5Loader

        loader = HDF5Loader()
        df = loader.load_episodes(hdf5_file)
        assert not df.empty
        assert "episode_index" in df.columns
        assert "action" in df.columns
        assert "state" in df.columns
        assert df["episode_index"].nunique() == 5

        # Check that actions have correct dimensionality
        first_action = df.iloc[0]["action"]
        assert len(first_action) == 7

    def test_load_specific_episodes(self, hdf5_file):
        from orbit.loaders.hdf5 import HDF5Loader

        loader = HDF5Loader()
        df = loader.load_episodes(hdf5_file, episode_indices=[0, 2, 4])
        assert df["episode_index"].nunique() == 3


# ---------------------------------------------------------------------------
# Directory Loader
# ---------------------------------------------------------------------------


class TestDirectoryLoader:
    def test_per_episode_npz(self, tmp_path):
        from orbit.loaders.directory import DirectoryLoader

        # Create episode NPZ files
        for i in range(3):
            n = 20 + i * 10
            np.savez(
                tmp_path / f"episode_{i:03d}.npz",
                action=np.random.randn(n, 6).astype(np.float32),
                state=np.random.randn(n, 6).astype(np.float32),
            )

        loader = DirectoryLoader()
        assert loader.can_load(str(tmp_path)) is True

        info = loader.load_info(str(tmp_path))
        assert info.num_episodes == 3
        assert info.shapes["action"] == [6]

        df = loader.load_episodes(str(tmp_path))
        assert not df.empty
        assert df["episode_index"].nunique() == 3

    def test_flat_arrays(self, tmp_path):
        from orbit.loaders.directory import DirectoryLoader

        # Create flat arrays + episode boundaries
        n_total = 100
        np.save(tmp_path / "actions.npy", np.random.randn(n_total, 7).astype(np.float32))
        np.save(tmp_path / "states.npy", np.random.randn(n_total, 14).astype(np.float32))
        with open(tmp_path / "episodes.json", "w") as f:
            json.dump(
                [
                    {"start": 0, "end": 40},
                    {"start": 40, "end": 70},
                    {"start": 70, "end": 100},
                ],
                f,
            )

        loader = DirectoryLoader()
        assert loader.can_load(str(tmp_path)) is True

        info = loader.load_info(str(tmp_path))
        assert info.num_episodes == 3
        assert info.num_frames == 100

        df = loader.load_episodes(str(tmp_path))
        assert not df.empty
        assert df["episode_index"].nunique() == 3
        # Check frame counts per episode
        counts = df.groupby("episode_index").size()
        assert counts[0] == 40
        assert counts[1] == 30
        assert counts[2] == 30

    def test_tabular_parquet(self, tmp_path):
        from orbit.loaders.directory import DirectoryLoader

        # Create a parquet with action/state columns
        df = pd.DataFrame(
            {
                "episode_index": [0] * 10 + [1] * 10,
                "frame_index": list(range(10)) * 2,
                "action_0": np.random.randn(20),
                "action_1": np.random.randn(20),
                "state_0": np.random.randn(20),
                "state_1": np.random.randn(20),
                "timestamp": np.arange(20) * 0.05,
            }
        )
        df.to_parquet(tmp_path / "data.parquet", index=False)

        loader = DirectoryLoader()
        assert loader.can_load(str(tmp_path)) is True

        info = loader.load_info(str(tmp_path))
        assert info.num_episodes == 2

        result = loader.load_episodes(str(tmp_path))
        assert not result.empty
        assert result["episode_index"].nunique() == 2
        # Actions should be combined into a list
        first_action = result.iloc[0]["action"]
        assert isinstance(first_action, list)
        assert len(first_action) == 2

    def test_does_not_match_lerobot_dir(self, tmp_path):
        from orbit.loaders.directory import DirectoryLoader

        meta = tmp_path / "meta"
        meta.mkdir()
        (meta / "info.json").write_text("{}")
        np.savez(tmp_path / "episode_000.npz", action=np.zeros((5, 3)))

        loader = DirectoryLoader()
        # Should NOT match because meta/info.json exists (LeRobot)
        assert loader.can_load(str(tmp_path)) is False


# ---------------------------------------------------------------------------
# ROS Bag Loader (unit tests for synchronization/segmentation logic)
# ---------------------------------------------------------------------------


class TestRosBagLoaderUnit:
    def test_segment_episodes(self):
        from orbit.loaders.rosbag import RosBagLoader

        loader = RosBagLoader(episode_gap=1.0)

        raw = {
            "timestamps": [0.0, 0.1, 0.2, 0.3, 2.0, 2.1, 2.2],
            "actions": [[1.0]] * 7,
            "states": [[2.0]] * 7,
        }

        episodes = loader._segment_episodes(raw)
        assert len(episodes) == 2
        assert len(episodes[0]["timestamps"]) == 4
        assert len(episodes[1]["timestamps"]) == 3

    def test_synchronize_messages(self):
        from orbit.loaders.rosbag import RosBagLoader

        loader = RosBagLoader()

        action_msgs = [(0.0, [1.0, 2.0]), (0.1, [3.0, 4.0]), (0.2, [5.0, 6.0])]
        state_msgs = [(0.05, [10.0, 20.0]), (0.15, [30.0, 40.0])]

        result = loader._synchronize_messages(action_msgs, state_msgs)
        assert len(result["timestamps"]) == 3
        assert len(result["actions"]) == 3
        assert len(result["states"]) == 3
        # First action at t=0.0 should pair with first state at t=0.05
        assert result["states"][0] == [10.0, 20.0]

    def test_can_load_mcap(self, tmp_path):
        from orbit.loaders.rosbag import RosBagLoader

        loader = RosBagLoader()
        # Just extension check
        p = tmp_path / "test.mcap"
        p.write_bytes(b"")
        assert loader.can_load(str(p)) is True

    def test_can_load_bag(self, tmp_path):
        from orbit.loaders.rosbag import RosBagLoader

        loader = RosBagLoader()
        p = tmp_path / "test.bag"
        p.write_bytes(b"")
        assert loader.can_load(str(p)) is True

    def test_topic_detection(self):
        from orbit.loaders.rosbag import RosBagLoader

        loader = RosBagLoader()
        topics = ["/camera/image", "/joint_states", "/joint_commands", "/rosout"]

        action = loader._find_topic(topics, loader.ACTION_TOPIC_PATTERNS)
        state = loader._find_topic(topics, loader.STATE_TOPIC_PATTERNS)

        assert action == "/joint_commands"
        assert state == "/joint_states"


# ---------------------------------------------------------------------------
# RLDS Loader (basic unit tests)
# ---------------------------------------------------------------------------


class TestRLDSLoaderUnit:
    def test_can_load_tfrecord_dir(self, tmp_path):
        from orbit.loaders.rlds import RLDSLoader

        loader = RLDSLoader()
        # Empty dir
        assert loader.can_load(str(tmp_path)) is False

        # Dir with tfrecord file
        (tmp_path / "data.tfrecord-00000-of-00001").write_bytes(b"")
        assert loader.can_load(str(tmp_path)) is True

    def test_can_load_tfds_dir(self, tmp_path):
        from orbit.loaders.rlds import RLDSLoader

        loader = RLDSLoader()
        (tmp_path / "dataset_info.json").write_text("{}")
        assert loader.can_load(str(tmp_path)) is True

    def test_load_info_from_metadata(self, tmp_path):
        from orbit.loaders.rlds import RLDSLoader

        loader = RLDSLoader()
        (tmp_path / "dataset_info.json").write_text(
            json.dumps(
                {
                    "splits": [{"statistics": {"numExamples": 100}}],
                    "features": {
                        "steps": {
                            "feature": {
                                "action": {"shape": {"dimensions": ["7"]}},
                                "observation": {
                                    "state": {"shape": {"dimensions": ["14"]}}
                                },
                            }
                        }
                    },
                }
            )
        )
        # Create a dummy tfrecord so can_load works
        (tmp_path / "dummy.tfrecord").write_bytes(b"")

        info = loader.load_info(str(tmp_path))
        assert info.num_episodes == 100
        assert info.shapes.get("action") == [7]
        assert info.shapes.get("state") == [14]


# ---------------------------------------------------------------------------
# Convert
# ---------------------------------------------------------------------------


class TestConvert:
    def test_write_lerobot_v3(self, tmp_path):
        from orbit.loaders.convert import write_lerobot_v3

        info = DatasetInfo(
            repo_id="test",
            num_episodes=2,
            num_frames=20,
            fps=30.0,
            robot_type="so100",
            camera_names=[],
            video_keys=[],
            shapes={"action": [6], "state": [6]},
        )

        df = pd.DataFrame(
            {
                "episode_index": [0] * 10 + [1] * 10,
                "frame_index": list(range(10)) * 2,
                "action": [np.random.randn(6).tolist() for _ in range(20)],
                "state": [np.random.randn(6).tolist() for _ in range(20)],
                "timestamp": [i / 30.0 for i in range(10)] * 2,
            }
        )

        output = tmp_path / "output"
        write_lerobot_v3(info, df, output)

        # Verify structure
        assert (output / "meta" / "info.json").exists()
        assert (output / "meta" / "episodes" / "chunk-000" / "file-000.parquet").exists()
        assert (output / "data" / "chunk-000" / "episode_000000.parquet").exists()
        assert (output / "data" / "chunk-000" / "episode_000001.parquet").exists()

        # Verify info.json
        with open(output / "meta" / "info.json") as f:
            out_info = json.load(f)
        assert out_info["total_episodes"] == 2
        assert out_info["total_frames"] == 20
        assert out_info["fps"] == 30.0

        # Verify episode data is readable
        ep0 = pd.read_parquet(output / "data" / "chunk-000" / "episode_000000.parquet")
        assert len(ep0) == 10
        assert "action" in ep0.columns


# ---------------------------------------------------------------------------
# Backward compatibility
# ---------------------------------------------------------------------------


class TestBackwardCompat:
    """Verify the refactored dataset_loader.py still exports the same symbols."""

    def test_imports(self):
        from orbit.analyzer.dataset_loader import (
            DatasetInfo,
            _sample_episode_indices,
            load_dataset_info,
            load_episode_data,
        )

        assert callable(load_dataset_info)
        assert callable(load_episode_data)
        assert callable(_sample_episode_indices)
        assert DatasetInfo is not None

    def test_sample_episode_indices(self):
        from orbit.analyzer.dataset_loader import _sample_episode_indices

        assert _sample_episode_indices(5, 10) == [0, 1, 2, 3, 4]
        result = _sample_episode_indices(100, 10)
        assert len(result) == 10

    def test_load_info_dispatches_to_local(self, tmp_path):
        """Verify load_dataset_info works with local LeRobot dirs."""
        from orbit.analyzer.dataset_loader import load_dataset_info

        meta = tmp_path / "meta"
        meta.mkdir()
        (meta / "info.json").write_text(
            json.dumps(
                {
                    "total_episodes": 5,
                    "total_frames": 250,
                    "fps": 50.0,
                    "robot_type": "koch",
                    "features": {},
                }
            )
        )

        info = load_dataset_info(str(tmp_path))
        assert info.num_episodes == 5
        assert info.robot_type == "koch"


# ---------------------------------------------------------------------------
# End-to-end: directory loader through the dispatcher
# ---------------------------------------------------------------------------


class TestEndToEnd:
    def test_npz_through_dispatcher(self, tmp_path):
        from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data

        for i in range(2):
            np.savez(
                tmp_path / f"episode_{i:03d}.npz",
                action=np.random.randn(10, 4).astype(np.float32),
                state=np.random.randn(10, 4).astype(np.float32),
            )

        info = load_dataset_info(str(tmp_path))
        assert info.num_episodes == 2

        df = load_episode_data(str(tmp_path))
        assert not df.empty
        assert "action" in df.columns
        assert "state" in df.columns
        assert df["episode_index"].nunique() == 2

    def test_hdf5_through_dispatcher(self, tmp_path):
        h5py = pytest.importorskip("h5py")
        from orbit.analyzer.dataset_loader import load_dataset_info, load_episode_data

        p = tmp_path / "test.hdf5"
        with h5py.File(str(p), "w") as f:
            data = f.create_group("data")
            for i in range(3):
                demo = data.create_group(f"demo_{i}")
                demo.create_dataset("actions", data=np.random.randn(15, 7))
                demo.create_dataset("states", data=np.random.randn(15, 14))

        info = load_dataset_info(str(p))
        assert info.num_episodes == 3

        df = load_episode_data(str(p))
        assert not df.empty
        assert df["episode_index"].nunique() == 3
