"""Tests for orbit.analyzer.dataset_loader."""

from __future__ import annotations

import json
from unittest.mock import MagicMock, patch

import pandas as pd
import pytest

from orbit.analyzer.dataset_loader import (
    DatasetInfo,
    _sample_episode_indices,
    load_dataset_info,
    load_episode_data,
)

# ---------------------------------------------------------------------------
# Fixtures / helpers
# ---------------------------------------------------------------------------

SAMPLE_INFO_JSON = {
    "total_episodes": 50,
    "total_frames": 25000,
    "fps": 50.0,
    "robot_type": "aloha",
    "features": {
        "observation.images.top": {"dtype": "video", "shape": [480, 640, 3]},
        "observation.images.wrist": {"dtype": "video", "shape": [480, 640, 3]},
        "action": {"dtype": "float32", "shape": [14]},
        "observation.state": {"dtype": "float32", "shape": [14]},
    },
}

SAMPLE_EPISODES_JSONL = "\n".join(
    json.dumps({"episode_index": i, "length": 500}) for i in range(50)
)


@pytest.fixture()
def _mock_info_json(tmp_path):
    """Write a fake info.json and return its path."""
    p = tmp_path / "info.json"
    p.write_text(json.dumps(SAMPLE_INFO_JSON))
    return str(p)


@pytest.fixture()
def _mock_episodes_jsonl(tmp_path):
    """Write a fake episodes.jsonl and return its path."""
    p = tmp_path / "episodes.jsonl"
    p.write_text(SAMPLE_EPISODES_JSONL)
    return str(p)


# ---------------------------------------------------------------------------
# load_dataset_info
# ---------------------------------------------------------------------------


class TestLoadDatasetInfo:
    """Tests for load_dataset_info."""

    @patch("orbit.loaders.lerobot_hub.hf_hub_download")
    @patch("orbit.loaders.lerobot_hub.HfApi")
    def test_basic_metadata(
        self, mock_api_cls, mock_download, _mock_info_json, _mock_episodes_jsonl
    ):
        mock_api_cls.return_value.dataset_info.return_value = MagicMock()

        def _download(repo_id, filename, repo_type):
            if filename == "meta/info.json":
                return _mock_info_json
            if filename == "meta/episodes.jsonl":
                return _mock_episodes_jsonl
            raise FileNotFoundError(filename)

        mock_download.side_effect = _download

        result = load_dataset_info("lerobot/aloha_mobile_cabinet")

        assert isinstance(result, DatasetInfo)
        assert result.repo_id == "lerobot/aloha_mobile_cabinet"
        assert result.num_episodes == 50
        assert result.num_frames == 25000
        assert result.fps == 50.0
        assert result.robot_type == "aloha"
        assert result.duration_seconds == pytest.approx(500.0)
        assert sorted(result.camera_names) == ["top", "wrist"]
        assert sorted(result.video_keys) == [
            "observation.images.top",
            "observation.images.wrist",
        ]
        assert result.shapes["action"] == [14]
        assert result.shapes["observation.state"] == [14]

    @patch("orbit.loaders.lerobot_hub.hf_hub_download")
    def test_repo_not_found(self, mock_download):
        from unittest.mock import MagicMock

        from huggingface_hub.utils import RepositoryNotFoundError

        # RepositoryNotFoundError requires a response kwarg in newer versions
        try:
            exc = RepositoryNotFoundError("not found")
        except TypeError:
            mock_response = MagicMock()
            mock_response.status_code = 404
            mock_response.headers = {}
            exc = RepositoryNotFoundError("not found", response=mock_response)

        mock_download.side_effect = exc

        with pytest.raises(ValueError, match="not found on HuggingFace Hub"):
            load_dataset_info("fake/nonexistent")

    @patch("orbit.loaders.lerobot_hub.hf_hub_download")
    @patch("orbit.loaders.lerobot_hub.HfApi")
    def test_not_lerobot_dataset(self, mock_api_cls, mock_download):
        from huggingface_hub.utils import EntryNotFoundError

        mock_api_cls.return_value.dataset_info.return_value = MagicMock()
        mock_download.side_effect = EntryNotFoundError("meta/info.json not found")

        with pytest.raises(
            ValueError, match="does not appear to be a LeRobot dataset"
        ):
            load_dataset_info("some/text-dataset")


# ---------------------------------------------------------------------------
# Smart sampling
# ---------------------------------------------------------------------------


class TestSampling:
    """Tests for _sample_episode_indices."""

    def test_small_dataset_returns_all(self):
        indices = _sample_episode_indices(50, max_samples=100)
        assert indices == list(range(50))

    def test_exact_boundary_returns_all(self):
        indices = _sample_episode_indices(100, max_samples=100)
        assert indices == list(range(100))

    def test_large_dataset_samples_evenly(self):
        indices = _sample_episode_indices(1000, max_samples=100)
        assert len(indices) == 100
        assert indices[0] == 0
        # Last index should be near the end
        assert indices[-1] == 990
        # Check spacing is roughly uniform
        diffs = [indices[i + 1] - indices[i] for i in range(len(indices) - 1)]
        assert all(d == 10 for d in diffs)

    def test_odd_total(self):
        indices = _sample_episode_indices(257, max_samples=100)
        assert len(indices) == 100
        assert indices[0] == 0
        assert all(0 <= i < 257 for i in indices)


# ---------------------------------------------------------------------------
# load_episode_data
# ---------------------------------------------------------------------------


class TestLoadEpisodeData:
    """Tests for load_episode_data."""

    @patch("orbit.loaders.lerobot_hub.hf_hub_download")
    @patch("orbit.loaders.lerobot_hub.HfApi")
    def test_loads_parquet_files(
        self, mock_api_cls, mock_download, tmp_path, _mock_info_json, _mock_episodes_jsonl
    ):
        mock_api_cls.return_value.dataset_info.return_value = MagicMock()
        # Mock list_repo_files for v1 format discovery
        mock_api_cls.return_value.list_repo_files.return_value = [
            "meta/info.json",
            "meta/episodes.jsonl",
            "data/episode_000000.parquet",
        ]

        # Create a small parquet file for episode 0
        ep_df = pd.DataFrame(
            {
                "episode_index": [0, 0, 0],
                "frame_index": [0, 1, 2],
                "action": [[0.1] * 14, [0.2] * 14, [0.3] * 14],
                "state": [[1.0] * 14, [1.1] * 14, [1.2] * 14],
                "timestamp": [0.0, 0.02, 0.04],
            }
        )
        parquet_path = tmp_path / "episode_000000.parquet"
        ep_df.to_parquet(parquet_path)

        def _download(repo_id, filename, repo_type):
            if filename == "meta/info.json":
                return _mock_info_json
            if filename == "meta/episodes.jsonl":
                return _mock_episodes_jsonl
            if filename == "data/episode_000000.parquet":
                return str(parquet_path)
            from huggingface_hub.utils import EntryNotFoundError
            raise EntryNotFoundError(f"{filename} not found")

        mock_download.side_effect = _download

        result = load_episode_data(
            "lerobot/aloha_mobile_cabinet", episode_indices=[0]
        )

        assert isinstance(result, pd.DataFrame)
        assert len(result) == 3
        assert list(result["episode_index"]) == [0, 0, 0]
        assert list(result["frame_index"]) == [0, 1, 2]

    @patch("orbit.loaders.lerobot_hub.hf_hub_download")
    @patch("orbit.loaders.lerobot_hub.HfApi")
    def test_empty_result_for_missing_episodes(
        self, mock_api_cls, mock_download, _mock_info_json, _mock_episodes_jsonl
    ):
        mock_api_cls.return_value.dataset_info.return_value = MagicMock()

        def _download(repo_id, filename, repo_type):
            if filename == "meta/info.json":
                return _mock_info_json
            if filename == "meta/episodes.jsonl":
                return _mock_episodes_jsonl
            from huggingface_hub.utils import EntryNotFoundError
            raise EntryNotFoundError(f"{filename} not found")

        mock_download.side_effect = _download

        result = load_episode_data(
            "lerobot/aloha_mobile_cabinet", episode_indices=[999]
        )

        assert isinstance(result, pd.DataFrame)
        assert len(result) == 0
        assert "episode_index" in result.columns
