#!/usr/bin/env python

# Copyright 2024 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

import pytest
import torch
from packaging.version import Version

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

from datasets import Dataset  # noqa: E402
from huggingface_hub import DatasetCard

import lerobot.datasets.utils as dataset_utils
from lerobot.datasets.io_utils import hf_transform_to_torch
from lerobot.datasets.utils import (
    create_lerobot_dataset_card,
    get_repo_versions,
    get_safe_version,
    resolve_episode_indices,
)
from lerobot.utils.constants import ACTION, OBS_IMAGES
from lerobot.utils.feature_utils import combine_feature_dicts


def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
    """Calculate episode data index for testing. Returns {"from": Tensor, "to": Tensor}."""
    episode_data_index: dict[str, list[int]] = {"from": [], "to": []}
    current_episode = None
    if len(hf_dataset) == 0:
        return {"from": torch.tensor([]), "to": torch.tensor([])}
    for idx, episode_idx in enumerate(hf_dataset["episode_index"]):
        if episode_idx != current_episode:
            episode_data_index["from"].append(idx)
            if current_episode is not None:
                episode_data_index["to"].append(idx)
            current_episode = episode_idx
    episode_data_index["to"].append(idx + 1)
    return {k: torch.tensor(v) for k, v in episode_data_index.items()}


def test_default_parameters():
    card = create_lerobot_dataset_card()
    assert isinstance(card, DatasetCard)
    assert card.data.tags == ["LeRobot"]
    assert card.data.task_categories == ["robotics"]
    assert card.data.configs == [
        {
            "config_name": "default",
            "data_files": "data/*/*.parquet",
        }
    ]


def test_resolve_episode_indices_applies_allowlist_and_exclusions():
    assert resolve_episode_indices([4, 1, 3, 0], 5, [1, 4]) == [3, 0]


def test_resolve_episode_indices_preserves_none_without_filtering():
    assert resolve_episode_indices(None, 5) is None


def test_resolve_episode_indices_ignores_out_of_range_values(caplog):
    assert resolve_episode_indices([-1, 0, 3, 5], 4, [-2, 3, 8]) == [0]
    assert "Ignoring episode indices outside the dataset range [0, 4): [-1, 5]" in caplog.text
    assert "Ignoring excluded episode indices outside the dataset range [0, 4): [-2, 8]" in caplog.text


@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_get_repo_versions_forwards_token(monkeypatch, token):
    api = Mock()
    api.list_repo_refs.return_value = SimpleNamespace(
        branches=[SimpleNamespace(name="v3.0")],
        tags=[],
    )
    hf_api = Mock(return_value=api)
    monkeypatch.setattr(dataset_utils, "HfApi", hf_api)

    assert get_repo_versions("private/repo", token=token) == [Version("3.0")]
    hf_api.assert_called_once_with(token=token)
    api.list_repo_refs.assert_called_once_with("private/repo", repo_type="dataset")


@pytest.mark.parametrize("token", ["hf_test_token", True, False])
def test_get_safe_version_forwards_token(monkeypatch, token):
    get_versions = Mock(return_value=[Version("3.0")])
    monkeypatch.setattr(dataset_utils, "get_repo_versions", get_versions)

    assert get_safe_version("private/repo", "v3.0", token=token) == "v3.0"
    get_versions.assert_called_once_with("private/repo", token=token)


def test_with_tags():
    tags = ["tag1", "tag2"]
    card = create_lerobot_dataset_card(tags=tags)
    assert card.data.tags == ["LeRobot", "tag1", "tag2"]


def test_calculate_episode_data_index():
    dataset = Dataset.from_dict(
        {
            "timestamp": [0.1, 0.2, 0.3, 0.4, 0.5, 0.6],
            "index": [0, 1, 2, 3, 4, 5],
            "episode_index": [0, 0, 1, 2, 2, 2],
        },
    )
    dataset.set_transform(hf_transform_to_torch)
    episode_data_index = calculate_episode_data_index(dataset)
    assert torch.equal(episode_data_index["from"], torch.tensor([0, 2, 3]))
    assert torch.equal(episode_data_index["to"], torch.tensor([2, 3, 6]))


def test_merge_simple_vectors():
    g1 = {
        ACTION: {
            "dtype": "float32",
            "shape": (2,),
            "names": ["ee.x", "ee.y"],
        }
    }
    g2 = {
        ACTION: {
            "dtype": "float32",
            "shape": (2,),
            "names": ["ee.y", "ee.z"],
        }
    }

    out = combine_feature_dicts(g1, g2)

    assert ACTION in out
    assert out[ACTION]["dtype"] == "float32"
    # Names merged with preserved order and de-dupuplication
    assert out[ACTION]["names"] == ["ee.x", "ee.y", "ee.z"]
    # Shape correctly recomputed from names length
    assert out[ACTION]["shape"] == (3,)


def test_merge_multiple_groups_order_and_dedup():
    g1 = {ACTION: {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}}
    g2 = {ACTION: {"dtype": "float32", "shape": (2,), "names": ["b", "c"]}}
    g3 = {ACTION: {"dtype": "float32", "shape": (3,), "names": ["a", "c", "d"]}}

    out = combine_feature_dicts(g1, g2, g3)

    assert out[ACTION]["names"] == ["a", "b", "c", "d"]
    assert out[ACTION]["shape"] == (4,)


def test_non_vector_last_wins_for_images():
    # Non-vector (images) with same name should be overwritten by the last image specified
    g1 = {
        f"{OBS_IMAGES}.front": {
            "dtype": "image",
            "shape": (3, 480, 640),
            "names": ["channels", "height", "width"],
        }
    }
    g2 = {
        f"{OBS_IMAGES}.front": {
            "dtype": "image",
            "shape": (3, 720, 1280),
            "names": ["channels", "height", "width"],
        }
    }

    out = combine_feature_dicts(g1, g2)
    assert out[f"{OBS_IMAGES}.front"]["shape"] == (3, 720, 1280)
    assert out[f"{OBS_IMAGES}.front"]["dtype"] == "image"


def test_dtype_mismatch_raises():
    g1 = {ACTION: {"dtype": "float32", "shape": (1,), "names": ["a"]}}
    g2 = {ACTION: {"dtype": "float64", "shape": (1,), "names": ["b"]}}

    with pytest.raises(ValueError, match="dtype mismatch for 'action'"):
        _ = combine_feature_dicts(g1, g2)


def test_non_dict_passthrough_last_wins():
    g1 = {"misc": 123}
    g2 = {"misc": 456}

    out = combine_feature_dicts(g1, g2)
    # For non-dict entries the last one wins
    assert out["misc"] == 456


def test_get_safe_version_raises_on_repo_without_version_tags(monkeypatch):
    monkeypatch.setattr(dataset_utils, "get_repo_versions", Mock(return_value=[]))

    with pytest.raises(RuntimeError, match="must be tagged with a codebase version"):
        get_safe_version("private/repo", "v3.0")


def test_get_safe_version_error_reports_repo_id(monkeypatch):
    repo_id = "private/repo"
    monkeypatch.setattr(dataset_utils, "get_repo_versions", Mock(return_value=[]))

    with pytest.raises(RuntimeError) as exc_info:
        get_safe_version(repo_id, "v3.0")

    assert repo_id in str(exc_info.value)


@pytest.mark.parametrize(
    ("anchor", "start", "end", "indices", "padding"),
    [
        # Absolute dataset indices (map-style): episode rows 10..14.
        (10, 10, 15, [10, 10, 11], [True, False, False]),
        (14, 10, 15, [13, 14, 14], [False, False, True]),
        # Episode-local indices (streaming): rows 0..4.
        (2, 0, 5, [1, 2, 3], [False, False, False]),
    ],
)
def test_delta_window_clamps_and_marks_padding(anchor, start, end, indices, padding):
    window, is_pad = dataset_utils.delta_window(anchor, [-1, 0, 1], start, end)
    assert window == indices
    assert is_pad.tolist() == padding
    assert is_pad.dtype == torch.bool


def test_shift_timestamps_and_task_name():
    import pandas as pd

    assert dataset_utils.shift_timestamps([0.0, 0.5], 10.0) == [10.0, 10.5]
    tasks = pd.DataFrame({"task_index": [0, 1]}, index=["pick", "place"])
    assert dataset_utils.task_name(tasks, torch.tensor(1)) == "place"
