#!/usr/bin/env python

# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from types import SimpleNamespace
from unittest.mock import Mock, patch

import numpy as np
import pytest
import torch

pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")

import lerobot.datasets.streaming_dataset as streaming_dataset_module
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
from lerobot.utils.constants import ACTION
from tests.fixtures.constants import DUMMY_REPO_ID

# A dataset whose videos roll over into a second file, as v3.0 does past
# DEFAULT_VIDEO_FILE_SIZE_IN_MB: episodes 4-7 live in file-001, whose timeline restarts at 0.
MULTI_FILE_EPISODES = 8
MULTI_FILE_FRAMES = 200
MULTI_FILE_EPISODES_PER_VIDEO_FILE = 4


def get_frames_expected_order(streaming_ds: StreamingLeRobotDataset) -> list[int]:
    """Replicates the shuffling logic of StreamingLeRobotDataset to get the expected order of indices."""
    rng = np.random.default_rng(streaming_ds.seed)
    buffer_size = streaming_ds.buffer_size
    num_shards = streaming_ds.num_shards

    shards_indices = []
    for shard_idx in range(num_shards):
        shard = streaming_ds.hf_dataset.shard(num_shards, index=shard_idx)
        shard_indices = [item["index"] for item in shard]
        shards_indices.append(shard_indices)

    shard_iterators = {i: iter(s) for i, s in enumerate(shards_indices)}

    buffer_indices_generator = streaming_ds._iter_random_indices(rng, buffer_size)

    frames_buffer = []
    expected_indices = []

    while shard_iterators:  # While there are still available shards
        available_shard_keys = list(shard_iterators.keys())
        if not available_shard_keys:
            break

        # Call _infinite_generator_over_elements with current available shards (key difference!)
        shard_key = next(streaming_ds._infinite_generator_over_elements(rng, available_shard_keys))

        try:
            frame_index = next(shard_iterators[shard_key])

            if len(frames_buffer) == buffer_size:
                i = next(buffer_indices_generator)
                expected_indices.append(frames_buffer[i])
                frames_buffer[i] = frame_index
            else:
                frames_buffer.append(frame_index)

        except StopIteration:
            del shard_iterators[shard_key]  # Remove exhausted shard

    rng.shuffle(frames_buffer)
    expected_indices.extend(frames_buffer)

    return expected_indices


@pytest.mark.parametrize("token", ["hf_test_token", True, False])
@pytest.mark.parametrize("from_local", [False, True])
def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, monkeypatch, token, from_local):
    requested_root = tmp_path / "local" if from_local else None
    metadata = _fake_meta(
        root=requested_root or tmp_path / "snapshot",
        revision=streaming_dataset_module.CODEBASE_VERSION,
    )
    metadata_cls = Mock(return_value=metadata)
    load_dataset = Mock(return_value=SimpleNamespace(num_shards=1))
    monkeypatch.setattr(streaming_dataset_module, "LeRobotDatasetMetadata", metadata_cls)
    monkeypatch.setattr(streaming_dataset_module, "load_dataset", load_dataset)

    dataset = StreamingLeRobotDataset(DUMMY_REPO_ID, root=requested_root, token=token)

    metadata_cls.assert_called_once_with(
        DUMMY_REPO_ID,
        requested_root,
        streaming_dataset_module.CODEBASE_VERSION,
        force_cache_sync=False,
        repo_type="dataset",
        token=token,
    )
    if from_local:
        assert "token" not in load_dataset.call_args.kwargs
    else:
        assert load_dataset.call_args.kwargs["token"] is token
    assert not hasattr(dataset, "_token")


def assert_videos_roll_over(ds: LeRobotDataset) -> None:
    """Videos spanning several files must decode each frame from its own file (v3.0 rollover)."""
    for key in ds.meta.video_keys:
        episodes = [ds.meta.episodes[ep_idx] for ep_idx in range(ds.meta.total_episodes)]
        file_indices = {ep[f"videos/{key}/file_index"] for ep in episodes}
        assert len(file_indices) > 1, f"{key} is not split across video files (file_index: {file_indices})"

        # The property under test: a later file's timeline starts back at 0, so an episode's
        # `from_timestamp` is no longer its global position in the dataset.
        assert any(
            ep[f"videos/{key}/file_index"] > 0 and ep[f"videos/{key}/from_timestamp"] == 0.0
            for ep in episodes
        ), f"No episode of {key} restarts a video file's timeline at 0"


def assert_frame_matches(streaming_frame: dict, target_frame: dict, ds: LeRobotDataset, context: str) -> None:
    """Assert a streamed frame equals the same frame read by the non-streaming reader."""
    assert set(streaming_frame.keys()) == set(target_frame.keys()), (
        f"Keys differ between streaming frame and target one ({context}). "
        f"Differ at: {set(streaming_frame.keys()) ^ set(target_frame.keys())}"
    )

    mismatched = []
    for key in streaming_frame:
        left, right = streaming_frame[key], target_frame[key]

        if isinstance(left, str):
            check = left == right

        elif isinstance(left, float):
            check = left == right.item()  # right is a torch.Tensor

        elif isinstance(left, torch.Tensor):
            if key not in ds.meta.camera_keys and "is_pad" not in key and f"{key}_is_pad" in streaming_frame:
                # comparing frames only on non-padded regions. Padding is applied to last-valid broadcasting
                left = left[~streaming_frame[f"{key}_is_pad"]]
                right = right[~target_frame[f"{key}_is_pad"]]

            check = left.shape == right.shape and torch.allclose(left, right)

        else:
            check = left == right

        if not check:
            mismatched.append(key)

    assert not mismatched, f"Streaming and target frames differ on {mismatched} ({context})"


def assert_stream_matches_reference(
    streaming_ds: StreamingLeRobotDataset, ds: LeRobotDataset, num_frames: int
) -> None:
    """Stream ``num_frames`` frames and assert each equals the same frame from the non-streaming reader."""
    stream = iter(streaming_ds)
    for i in range(num_frames):
        streaming_frame = next(stream)
        frame_idx = streaming_frame["index"]
        assert_frame_matches(streaming_frame, ds[frame_idx], ds, context=f"i: {i}, frame_idx: {frame_idx}")


def test_single_frame_consistency(tmp_path, lerobot_dataset_factory):
    """Streaming without deltas returns the same frames as the non-streaming reader."""
    ds_num_frames = 400
    ds_num_episodes = 10
    buffer_size = 100

    local_path = tmp_path / "test"
    repo_id = f"{DUMMY_REPO_ID}"

    ds = lerobot_dataset_factory(
        root=local_path,
        repo_id=repo_id,
        total_episodes=ds_num_episodes,
        total_frames=ds_num_frames,
    )

    streaming_ds = StreamingLeRobotDataset(repo_id=repo_id, root=local_path, buffer_size=buffer_size)
    assert_stream_matches_reference(streaming_ds, ds, ds_num_frames)


@pytest.mark.parametrize(
    "shuffle",
    [False, True],
)
def test_frames_order_over_epochs(tmp_path, lerobot_dataset_factory, shuffle):
    """Test if streamed frames correspond to shuffling operations over in-memory dataset."""
    ds_num_frames = 400
    ds_num_episodes = 10
    buffer_size = 100
    seed = 42
    n_epochs = 3

    local_path = tmp_path / "test"
    repo_id = f"{DUMMY_REPO_ID}"

    lerobot_dataset_factory(
        root=local_path,
        repo_id=repo_id,
        total_episodes=ds_num_episodes,
        total_frames=ds_num_frames,
    )

    streaming_ds = StreamingLeRobotDataset(
        repo_id=repo_id, root=local_path, buffer_size=buffer_size, seed=seed, shuffle=shuffle
    )

    first_epoch_indices = [frame["index"] for frame in streaming_ds]
    expected_indices = get_frames_expected_order(streaming_ds)

    assert first_epoch_indices == expected_indices, "First epoch indices do not match expected indices"

    expected_indices = get_frames_expected_order(streaming_ds)
    for _ in range(n_epochs):
        streaming_indices = [frame["index"] for frame in streaming_ds]
        frames_match = all(
            s_index == e_index for s_index, e_index in zip(streaming_indices, expected_indices, strict=True)
        )

        if shuffle:
            assert not frames_match
        else:
            assert frames_match


@pytest.mark.parametrize(
    "shuffle",
    [False, True],
)
def test_frames_order_with_shards(tmp_path, lerobot_dataset_factory, shuffle):
    """Test if streamed frames correspond to shuffling operations over in-memory dataset with multiple shards."""
    ds_num_frames = 100
    ds_num_episodes = 10
    buffer_size = 10

    seed = 42
    n_epochs = 3
    data_file_size_mb = 0.001

    chunks_size = 1

    local_path = tmp_path / "test"
    repo_id = f"{DUMMY_REPO_ID}-ciao"

    lerobot_dataset_factory(
        root=local_path,
        repo_id=repo_id,
        total_episodes=ds_num_episodes,
        total_frames=ds_num_frames,
        data_files_size_in_mb=data_file_size_mb,
        chunks_size=chunks_size,
    )

    streaming_ds = StreamingLeRobotDataset(
        repo_id=repo_id,
        root=local_path,
        buffer_size=buffer_size,
        seed=seed,
        shuffle=shuffle,
        max_num_shards=4,
    )

    first_epoch_indices = [frame["index"] for frame in streaming_ds]
    expected_indices = get_frames_expected_order(streaming_ds)

    assert first_epoch_indices == expected_indices, "First epoch indices do not match expected indices"

    for _ in range(n_epochs):
        streaming_indices = [
            frame["index"] for frame in streaming_ds
        ]  # NOTE: this is the same as first_epoch_indices
        frames_match = all(
            s_index == e_index for s_index, e_index in zip(streaming_indices, expected_indices, strict=True)
        )
        if shuffle:
            assert not frames_match
        else:
            assert frames_match


def test_iter_raises_on_frame_error(tmp_path, lerobot_dataset_factory, monkeypatch):
    """Video decode failures must propagate instead of being silently ignored."""
    ds_num_frames = 20
    ds_num_episodes = 2
    buffer_size = 10

    local_path = tmp_path / "test"
    repo_id = f"{DUMMY_REPO_ID}"

    lerobot_dataset_factory(
        root=local_path,
        repo_id=repo_id,
        total_episodes=ds_num_episodes,
        total_frames=ds_num_frames,
    )

    streaming_ds = StreamingLeRobotDataset(repo_id=repo_id, root=local_path, buffer_size=buffer_size)

    def broken_video_decode(*args, **kwargs):
        raise RuntimeError("Could not load libtorchcodec")

    monkeypatch.setattr(streaming_dataset_module, "decode_video_frames_torchcodec", broken_video_decode)

    with pytest.raises(RuntimeError, match="libtorchcodec"):
        next(iter(streaming_ds))


def test_iter_raises_on_nested_generator_error(tmp_path, lerobot_dataset_factory, monkeypatch):
    """PEP 479 errors below frame construction must not be mistaken for shard exhaustion."""
    local_path = tmp_path / "test"
    repo_id = DUMMY_REPO_ID

    lerobot_dataset_factory(root=local_path, repo_id=repo_id, total_episodes=2, total_frames=20)
    streaming_ds = StreamingLeRobotDataset(repo_id=repo_id, root=local_path, buffer_size=10)

    def broken_frame_generator():
        raise StopIteration("decoder internal failure")
        yield

    def broken_video_decode(*args, **kwargs):
        return next(broken_frame_generator())

    monkeypatch.setattr(streaming_dataset_module, "decode_video_frames_torchcodec", broken_video_decode)

    with pytest.raises(RuntimeError, match="generator raised StopIteration"):
        next(iter(streaming_ds))


@pytest.mark.parametrize("sharded", [False, True])
@pytest.mark.parametrize(
    "state_deltas, action_deltas",
    [
        ([-1, -0.5, -0.20, 0], [0, 1, 2, 3, 10, 20]),
        ([-1, -0.5, -0.20, 0], [-20, -1.5, -1, -0.5, -0.20, -0.10, 0]),
        ([-2, -1, -0.5, 0], [0, 1, 2, 3, 10, 20]),
        ([-2, -1, -0.5, 0], [-20, -1.5, -1, -0.5, -0.20, -0.10, 0]),
    ],
)
def test_frames_with_delta_consistency(
    tmp_path, lerobot_dataset_factory, sharded, state_deltas, action_deltas
):
    """Delta-window frames streamed match the non-streaming reader, with and without sharding."""
    local_path = tmp_path / "test"
    repo_id = f"{DUMMY_REPO_ID}-ciao"
    delta_timestamps = {"phone": state_deltas, "state": state_deltas, ACTION: action_deltas}

    if sharded:
        num_frames, buffer_size = 100, 10
        factory_extra = {"data_files_size_in_mb": 0.001, "chunks_size": 1}
        stream_extra = {"max_num_shards": 4}
    else:
        num_frames, buffer_size = 500, 100
        factory_extra, stream_extra = {}, {}

    ds = lerobot_dataset_factory(
        root=local_path,
        repo_id=repo_id,
        total_episodes=10,
        total_frames=num_frames,
        delta_timestamps=delta_timestamps,
        **factory_extra,
    )
    streaming_ds = StreamingLeRobotDataset(
        repo_id=repo_id,
        root=local_path,
        buffer_size=buffer_size,
        seed=42,
        shuffle=False,
        delta_timestamps=delta_timestamps,
        **stream_extra,
    )
    assert_stream_matches_reference(streaming_ds, ds, num_frames)


class _StopConstructionError(Exception):
    """Sentinel raised from a patched load_dataset to halt __init__ after the branch under test."""


def _fake_meta(*args, **kwargs):
    """Minimal LeRobotDatasetMetadata stand-in exposing only what __init__ reads."""
    meta = type("_Meta", (), {})()
    root = kwargs.get("root", args[1] if len(args) > 1 else None)
    revision = kwargs.get("revision", args[2] if len(args) > 2 else None)
    meta.root = root or "/tmp/_streaming_meta"
    meta.revision = revision or "v0"
    meta._version = streaming_dataset_module.CODEBASE_VERSION
    meta.features = {}
    meta.video_keys = []
    meta.depth_keys = []
    meta.image_keys = []
    meta.rescale_depth_stats = lambda *_a, **_k: None
    meta.repo_type = kwargs.get("repo_type", "dataset")
    return meta


@pytest.mark.parametrize(
    "repo_type, expected_source, expected_data_files",
    [
        ("bucket", "parquet", "hf://buckets/{repo_id}/data/*/*.parquet"),
        ("dataset", "{repo_id}", "data/*/*.parquet"),
    ],
)
def test_streaming_repo_type_routes_load_dataset(repo_type, expected_source, expected_data_files):
    """repo_type='bucket' loads parquet from hf://buckets/...; 'dataset' keeps the Hub-repo path."""
    captured = {}
    token = "hf_test_token"

    def fake_load_dataset(source, **kwargs):
        captured["source"] = source
        captured["data_files"] = kwargs.get("data_files")
        captured["token"] = kwargs.get("token")
        raise _StopConstructionError

    with (
        patch("lerobot.datasets.streaming_dataset.LeRobotDatasetMetadata", _fake_meta),
        patch("lerobot.datasets.streaming_dataset.check_version_compatibility", lambda *a, **k: None),
        patch("lerobot.datasets.streaming_dataset.load_dataset", fake_load_dataset),
        pytest.raises(_StopConstructionError),
    ):
        StreamingLeRobotDataset(DUMMY_REPO_ID, repo_type=repo_type, token=token)

    assert captured["source"] == expected_source.format(repo_id=DUMMY_REPO_ID)
    assert captured["data_files"] == expected_data_files.format(repo_id=DUMMY_REPO_ID)
    assert captured["token"] == token


def test_bucket_metadata_url_root(tmp_path):
    """repo_type='bucket' produces url_root pointing at hf://buckets/..."""
    with patch.object(LeRobotDatasetMetadata, "_load_metadata"):
        meta = LeRobotDatasetMetadata(
            DUMMY_REPO_ID,
            root=tmp_path,
            repo_type="bucket",
        )
    assert meta.url_root == f"hf://buckets/{DUMMY_REPO_ID}"


def test_bucket_skips_get_safe_version(tmp_path):
    """repo_type='bucket' must NOT call get_safe_version (buckets have no git refs).

    The first ``_load_metadata()`` raises ``FileNotFoundError`` so ``__init__``
    enters the except branch where the ``repo_type != "bucket"`` guard and
    ``get_safe_version`` live; the second call (after the bucket meta pull)
    succeeds. ``_pull_from_repo`` is stubbed so the except path runs without a
    network call. Without the raise, the try block would succeed and the guard
    branch would never execute, so the assertion would pass vacuously.
    """
    with (
        patch("lerobot.datasets.dataset_metadata.get_safe_version") as mock_gsv,
        patch.object(
            LeRobotDatasetMetadata,
            "_load_metadata",
            side_effect=[FileNotFoundError, None],
        ),
        patch.object(LeRobotDatasetMetadata, "_pull_from_repo") as mock_pull,
    ):
        LeRobotDatasetMetadata(
            DUMMY_REPO_ID,
            root=tmp_path,
            repo_type="bucket",
        )
    # The except branch ran (proven by the pull), but the bucket guard skipped
    # version resolution.
    mock_pull.assert_called_once()
    mock_gsv.assert_not_called()


def test_bucket_metadata_sync_uses_stable_cache_and_token(tmp_path):
    hub_cache = tmp_path / "hub"
    token = "hf_test_token"
    expected_root = hub_cache / f"buckets--{DUMMY_REPO_ID.replace('/', '--')}"

    with (
        patch("lerobot.datasets.dataset_metadata.HF_LEROBOT_HUB_CACHE", hub_cache),
        patch("lerobot.datasets.dataset_metadata.sync_bucket") as mock_sync,
        patch.object(
            LeRobotDatasetMetadata,
            "_load_metadata",
            side_effect=[FileNotFoundError, None],
        ),
    ):
        meta = LeRobotDatasetMetadata(DUMMY_REPO_ID, repo_type="bucket", token=token)

    assert meta.root == expected_root
    mock_sync.assert_called_once_with(
        f"hf://buckets/{DUMMY_REPO_ID}/meta",
        str(expected_root / "meta"),
        delete=True,
        quiet=True,
        token=token,
    )

    with (
        patch("lerobot.datasets.dataset_metadata.HF_LEROBOT_HUB_CACHE", hub_cache),
        patch("lerobot.datasets.dataset_metadata.sync_bucket") as mock_sync,
        patch.object(LeRobotDatasetMetadata, "_load_metadata"),
    ):
        cached_meta = LeRobotDatasetMetadata(DUMMY_REPO_ID, repo_type="bucket", token=token)

    assert cached_meta.root == expected_root
    mock_sync.assert_not_called()


def test_repo_type_is_keyword_only_and_preserves_positional_episodes():
    episodes = [1, 2]
    with (
        patch("lerobot.datasets.streaming_dataset.LeRobotDatasetMetadata", _fake_meta),
        patch("lerobot.datasets.streaming_dataset.check_version_compatibility"),
        patch(
            "lerobot.datasets.streaming_dataset.load_dataset",
            return_value=SimpleNamespace(num_shards=1),
        ),
    ):
        dataset = StreamingLeRobotDataset(DUMMY_REPO_ID, None, episodes)

    assert dataset.episodes == episodes
    assert dataset.repo_type == "dataset"


def test_bucket_root_caches_metadata_without_switching_to_local_streaming(tmp_path):
    with (
        patch("lerobot.datasets.streaming_dataset.LeRobotDatasetMetadata", _fake_meta),
        patch("lerobot.datasets.streaming_dataset.check_version_compatibility"),
        patch(
            "lerobot.datasets.streaming_dataset.load_dataset",
            return_value=SimpleNamespace(num_shards=1),
        ) as mock_load_dataset,
    ):
        dataset = StreamingLeRobotDataset(DUMMY_REPO_ID, root=tmp_path, repo_type="bucket")

    assert dataset.root == tmp_path
    assert not dataset.streaming_from_local
    assert mock_load_dataset.call_args.args == ("parquet",)
    assert mock_load_dataset.call_args.kwargs["data_files"].startswith("hf://buckets/")


def test_invalid_repo_type_fails_before_io():
    with pytest.raises(ValueError, match="repo_type must be 'dataset' or 'bucket'"):
        StreamingLeRobotDataset(DUMMY_REPO_ID, repo_type="space")


@pytest.mark.parametrize(
    "state_deltas, action_deltas",
    [
        (None, None),
        ([-1, -0.5, -0.20, 0], [0, 1, 2, 3]),
        ([-2, -1, -0.5, 0], [-1.5, -1, -0.5, -0.20, -0.10, 0]),
    ],
)
def test_consistency_across_video_files(tmp_path, lerobot_dataset_factory, state_deltas, action_deltas):
    """Videos spanning several files must decode each frame from its own file (v3.0 rollover)."""
    local_path = tmp_path / "test"
    repo_id = f"{DUMMY_REPO_ID}-video-rollover"
    delta_timestamps = (
        None
        if state_deltas is None
        else {"phone": state_deltas, "state": state_deltas, ACTION: action_deltas}
    )

    ds = lerobot_dataset_factory(
        root=local_path,
        repo_id=repo_id,
        total_episodes=MULTI_FILE_EPISODES,
        total_frames=MULTI_FILE_FRAMES,
        episodes_per_video_file=MULTI_FILE_EPISODES_PER_VIDEO_FILE,
        delta_timestamps=delta_timestamps,
    )
    assert_videos_roll_over(ds)

    streaming_ds = StreamingLeRobotDataset(
        repo_id=repo_id,
        root=local_path,
        buffer_size=100,
        seed=42,
        shuffle=False,
        delta_timestamps=delta_timestamps,
    )
    assert_stream_matches_reference(streaming_ds, ds, MULTI_FILE_FRAMES)