# 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.

"""Minimal tests for the rollout module's public API."""

from __future__ import annotations

import dataclasses
from unittest.mock import MagicMock

import pytest
import torch

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

# ---------------------------------------------------------------------------
# Import smoke tests
# ---------------------------------------------------------------------------


def test_rollout_top_level_imports():
    import lerobot.rollout

    for name in lerobot.rollout.__all__:
        assert hasattr(lerobot.rollout, name), f"Missing export: {name}"


def test_inference_submodule_imports():
    import lerobot.rollout.inference

    for name in lerobot.rollout.inference.__all__:
        assert hasattr(lerobot.rollout.inference, name), f"Missing export: {name}"


def test_strategies_submodule_imports():
    import lerobot.rollout.strategies

    for name in lerobot.rollout.strategies.__all__:
        assert hasattr(lerobot.rollout.strategies, name), f"Missing export: {name}"


# ---------------------------------------------------------------------------
# Config tests
# ---------------------------------------------------------------------------


def test_strategy_config_types():
    from lerobot.rollout import (
        BaseStrategyConfig,
        DAggerStrategyConfig,
        EpisodicStrategyConfig,
        HighlightStrategyConfig,
        SentryStrategyConfig,
    )

    assert BaseStrategyConfig().type == "base"
    assert SentryStrategyConfig().type == "sentry"
    assert HighlightStrategyConfig().type == "highlight"
    assert DAggerStrategyConfig().type == "dagger"
    assert EpisodicStrategyConfig().type == "episodic"


def test_dagger_config_invalid_input_device():
    from lerobot.rollout import DAggerStrategyConfig

    with pytest.raises(ValueError, match="input_device must be 'keyboard' or 'pedal'"):
        DAggerStrategyConfig(input_device="joystick")


def test_dagger_config_defaults():
    from lerobot.rollout import DAggerStrategyConfig

    cfg = DAggerStrategyConfig()
    assert cfg.num_episodes is None
    assert cfg.record_autonomous is False
    assert cfg.input_device == "keyboard"


def test_inference_config_types():
    from lerobot.rollout import RTCInferenceConfig, SyncInferenceConfig

    assert SyncInferenceConfig().type == "sync"

    rtc = RTCInferenceConfig()
    assert rtc.type == "rtc"
    assert rtc.queue_threshold == 30
    assert rtc.rtc is not None


def test_sentry_config_defaults():
    from lerobot.rollout import SentryStrategyConfig

    cfg = SentryStrategyConfig()
    assert cfg.upload_every_n_episodes == 5
    assert cfg.target_video_file_size_mb is None


# ---------------------------------------------------------------------------
# RolloutRingBuffer
# ---------------------------------------------------------------------------


def test_ring_buffer_append_and_eviction():
    from lerobot.rollout.ring_buffer import RolloutRingBuffer

    buf = RolloutRingBuffer(max_seconds=0.5, max_memory_mb=100.0, fps=10.0)
    # max_frames = 5
    for i in range(8):
        buf.append({"val": i})
    assert len(buf) == 5


def test_ring_buffer_drain():
    from lerobot.rollout.ring_buffer import RolloutRingBuffer

    buf = RolloutRingBuffer(max_seconds=1.0, max_memory_mb=100.0, fps=10.0)
    for i in range(3):
        buf.append({"val": i})
    frames = buf.drain()
    assert len(frames) == 3
    assert len(buf) == 0
    assert buf.estimated_bytes == 0


def test_ring_buffer_clear():
    from lerobot.rollout.ring_buffer import RolloutRingBuffer

    buf = RolloutRingBuffer(max_seconds=1.0, max_memory_mb=100.0, fps=10.0)
    buf.append({"val": 1})
    buf.clear()
    assert len(buf) == 0
    assert buf.estimated_bytes == 0


def test_ring_buffer_tensor_bytes():
    from lerobot.rollout.ring_buffer import RolloutRingBuffer

    buf = RolloutRingBuffer(max_seconds=1.0, max_memory_mb=100.0, fps=10.0)
    t = torch.zeros(100, dtype=torch.float32)  # 400 bytes
    buf.append({"tensor": t})
    assert buf.estimated_bytes >= 400


# ---------------------------------------------------------------------------
# ThreadSafeRobot
# ---------------------------------------------------------------------------


def test_thread_safe_robot_delegates():
    from lerobot.rollout.robot_wrapper import ThreadSafeRobot
    from tests.mocks.mock_robot import MockRobot, MockRobotConfig

    robot = MockRobot(MockRobotConfig(n_motors=3))
    robot.connect()
    wrapper = ThreadSafeRobot(robot)

    obs = wrapper.get_observation()
    assert "motor_1.pos" in obs
    assert "motor_2.pos" in obs
    assert "motor_3.pos" in obs

    action = {"motor_1.pos": 0.0, "motor_2.pos": 1.0, "motor_3.pos": 2.0}
    result = wrapper.send_action(action)
    assert result == action

    robot.disconnect()


def test_thread_safe_robot_properties():
    from lerobot.rollout.robot_wrapper import ThreadSafeRobot
    from tests.mocks.mock_robot import MockRobot, MockRobotConfig

    robot = MockRobot(MockRobotConfig(n_motors=3))
    robot.connect()
    wrapper = ThreadSafeRobot(robot)

    assert wrapper.name == "mock_robot"
    assert "motor_1.pos" in wrapper.observation_features
    assert "motor_1.pos" in wrapper.action_features
    assert wrapper.is_connected is True
    assert wrapper.inner is robot

    robot.disconnect()


# ---------------------------------------------------------------------------
# Strategy factory
# ---------------------------------------------------------------------------


def test_create_strategy_dispatches():
    from lerobot.rollout import (
        BaseStrategy,
        BaseStrategyConfig,
        DAggerStrategy,
        DAggerStrategyConfig,
        EpisodicStrategy,
        EpisodicStrategyConfig,
        SentryStrategy,
        SentryStrategyConfig,
        create_strategy,
    )

    assert isinstance(create_strategy(BaseStrategyConfig()), BaseStrategy)
    assert isinstance(create_strategy(SentryStrategyConfig()), SentryStrategy)
    assert isinstance(create_strategy(DAggerStrategyConfig()), DAggerStrategy)
    assert isinstance(create_strategy(EpisodicStrategyConfig()), EpisodicStrategy)


def test_create_strategy_unknown_raises():
    from lerobot.rollout import create_strategy

    cfg = MagicMock()
    cfg.type = "bogus"
    with pytest.raises(ValueError, match="Unknown strategy type"):
        create_strategy(cfg)


# ---------------------------------------------------------------------------
# Inference factory
# ---------------------------------------------------------------------------


def test_create_inference_engine_sync():
    from lerobot.rollout import SyncInferenceConfig, SyncInferenceEngine, create_inference_engine

    engine = create_inference_engine(
        SyncInferenceConfig(),
        policy=MagicMock(),
        preprocessor=MagicMock(),
        postprocessor=MagicMock(),
        robot_wrapper=MagicMock(robot_type="mock"),
        hw_features={},
        dataset_features={},
        ordered_action_keys=["k"],
        task="test",
        fps=30.0,
        device="cpu",
    )
    assert isinstance(engine, SyncInferenceEngine)


# ---------------------------------------------------------------------------
# Observation feature-spec consistency (sync vs RTC)
#
# ``build_dataset_frame`` orders the ``observation.state`` vector by the ``names``
# list of the feature spec it is handed. Sync uses ``dataset_features`` (joint
# layout AFTER ``robot_observation_processor``); RTC must use the SAME spec, not the
# raw-hardware layout, or its state vector desyncs from sync — corrupting both the
# normalizer's per-joint stats and the relative-action anchor, which sends the arm
# to systematically wrong absolute targets.
# ---------------------------------------------------------------------------


def test_build_dataset_frame_state_layout_depends_on_feature_spec():
    """A reordered feature spec yields a reordered ``observation.state`` vector."""
    import numpy as np

    from lerobot.utils.feature_utils import build_dataset_frame, hw_to_dataset_features

    # Raw hardware joint order.
    obs_hw = {"shoulder.pos": float, "elbow.pos": float, "wrist.pos": float, "gripper.pos": float}
    hw_features = hw_to_dataset_features(obs_hw, "observation")

    # What a ``robot_observation_processor`` that reorders state keys would produce
    # (e.g. gripper moved to the front). Same keys, different order.
    dataset_features = {
        "observation.state": {
            "dtype": "float32",
            "shape": (4,),
            "names": ["gripper.pos", "shoulder.pos", "elbow.pos", "wrist.pos"],
        }
    }

    values = {"shoulder.pos": 10.0, "elbow.pos": 20.0, "wrist.pos": 30.0, "gripper.pos": 40.0}

    hw_state = build_dataset_frame(hw_features, values, prefix="observation")["observation.state"]
    ds_state = build_dataset_frame(dataset_features, values, prefix="observation")["observation.state"]

    np.testing.assert_array_equal(hw_state, [10.0, 20.0, 30.0, 40.0])
    np.testing.assert_array_equal(ds_state, [40.0, 10.0, 20.0, 30.0])
    # The two layouts genuinely differ: this is the sync/RTC divergence the fix removes.
    assert not np.array_equal(hw_state, ds_state)


def test_create_inference_engine_rtc_uses_dataset_features():
    """The RTC engine must build observations from ``dataset_features`` (sync's spec),
    not the raw-hardware ``hw_features`` — so its ``observation.state`` matches sync."""
    from lerobot.rollout import RTCInferenceConfig, RTCInferenceEngine, create_inference_engine

    dataset_features = {
        "observation.state": {
            "dtype": "float32",
            "shape": (4,),
            "names": ["gripper.pos", "shoulder.pos", "elbow.pos", "wrist.pos"],
        }
    }
    # Deliberately different order to prove the engine ignores it.
    hw_features = {
        "observation.state": {
            "dtype": "float32",
            "shape": (4,),
            "names": ["shoulder.pos", "elbow.pos", "wrist.pos", "gripper.pos"],
        }
    }

    engine = create_inference_engine(
        RTCInferenceConfig(),
        policy=MagicMock(),
        # No relative/normalizer steps => __init__ introspection stays trivial.
        preprocessor=MagicMock(steps=[]),
        postprocessor=MagicMock(steps=[]),
        robot_wrapper=MagicMock(robot_type="mock"),
        hw_features=hw_features,
        dataset_features=dataset_features,
        ordered_action_keys=["k"],
        task="test",
        fps=30.0,
        device="cpu",
    )
    assert isinstance(engine, RTCInferenceEngine)
    assert engine._obs_features is dataset_features
    assert engine._obs_features["observation.state"]["names"] == [
        "gripper.pos",
        "shoulder.pos",
        "elbow.pos",
        "wrist.pos",
    ]


def test_rtc_get_action_remaps_model_order_to_ordered_action_keys():
    """RTC must remap the model-order action vector to ``ordered_action_keys`` by NAME
    before returning — matching sync. Otherwise the strategy maps model outputs onto the
    wrong joints (a per-joint permutation) whenever the two orders differ."""
    from lerobot.rollout import RTCInferenceConfig, RTCInferenceEngine, create_inference_engine
    from lerobot.utils.constants import ACTION

    # The model emits actions in dataset order [a, b, c]; the robot wants [c, a, b].
    dataset_action_names = ["a.pos", "b.pos", "c.pos"]
    ordered_action_keys = ["c.pos", "a.pos", "b.pos"]
    dataset_features = {
        ACTION: {"dtype": "float32", "shape": (3,), "names": dataset_action_names},
    }

    engine = create_inference_engine(
        RTCInferenceConfig(),
        policy=MagicMock(),
        preprocessor=MagicMock(steps=[]),
        postprocessor=MagicMock(steps=[]),
        robot_wrapper=MagicMock(robot_type="mock"),
        hw_features={},
        dataset_features=dataset_features,
        ordered_action_keys=ordered_action_keys,
        task="test",
        fps=30.0,
        device="cpu",
    )
    assert isinstance(engine, RTCInferenceEngine)

    # Queue yields the model-order vector a=1, b=2, c=3.
    engine._action_queue = MagicMock()
    engine._action_queue.get.return_value = torch.tensor([1.0, 2.0, 3.0])

    out = engine.get_action(None)
    # Remapped by name to [c, a, b] = [3, 1, 2]; positional pass-through would give [1, 2, 3].
    torch.testing.assert_close(out, torch.tensor([3.0, 1.0, 2.0]))


def test_normalize_prev_actions_length_holds_last_action_not_zeros():
    """A short RTC prefix is padded by repeating the last action, never with zeros —
    zeros decode to the mean action and cause the intermittent chunk-seam spike."""
    from lerobot.rollout.inference.rtc import _normalize_prev_actions_length

    prev = torch.tensor([[1.0, -1.0], [2.0, -2.0]])  # 2 steps, dim 2

    # Pad up to 5: rows 2..4 must equal the last real row, not zeros.
    padded = _normalize_prev_actions_length(prev, target_steps=5)
    assert padded.shape == (5, 2)
    torch.testing.assert_close(padded[:2], prev)
    for i in range(2, 5):
        torch.testing.assert_close(padded[i], prev[-1])
    assert not torch.any(padded[2:] == 0.0), "pad rows must not be zeros"

    # Exact length: unchanged. Truncation: first `target_steps` rows.
    torch.testing.assert_close(_normalize_prev_actions_length(prev, 2), prev)
    torch.testing.assert_close(_normalize_prev_actions_length(prev, 1), prev[:1])


# ---------------------------------------------------------------------------
# Pure functions
# ---------------------------------------------------------------------------


def test_estimate_max_episode_seconds_no_video():
    from lerobot.rollout.strategies import estimate_max_episode_seconds

    assert estimate_max_episode_seconds({}, fps=30.0) == 300.0


def test_estimate_max_episode_seconds_with_video():
    from lerobot.rollout.strategies import estimate_max_episode_seconds

    features = {"cam": {"dtype": "video", "shape": (480, 640, 3)}}
    result = estimate_max_episode_seconds(features, fps=30.0)
    assert result > 0
    # With a real camera, duration should differ from the fallback
    assert result != 300.0


def test_safe_push_to_hub():
    from lerobot.rollout.strategies import safe_push_to_hub

    ds = MagicMock()
    ds.num_episodes = 0
    assert safe_push_to_hub(ds) is False
    ds.push_to_hub.assert_not_called()

    ds.num_episodes = 5
    assert safe_push_to_hub(ds, tags=["test"]) is True
    ds.push_to_hub.assert_called_once_with(tags=["test"], private=False)


# ---------------------------------------------------------------------------
# DAgger state machine
# ---------------------------------------------------------------------------


def test_dagger_full_transition_cycle():
    from lerobot.rollout.strategies import DAggerEvents, DAggerPhase

    events = DAggerEvents()
    assert events.phase == DAggerPhase.AUTONOMOUS

    # AUTONOMOUS -> PAUSED
    events.request_transition("pause_resume")
    old, new = events.consume_transition()
    assert (old, new) == (DAggerPhase.AUTONOMOUS, DAggerPhase.PAUSED)

    # PAUSED -> CORRECTING
    events.request_transition("correction")
    old, new = events.consume_transition()
    assert (old, new) == (DAggerPhase.PAUSED, DAggerPhase.CORRECTING)

    # CORRECTING -> PAUSED
    events.request_transition("correction")
    old, new = events.consume_transition()
    assert (old, new) == (DAggerPhase.CORRECTING, DAggerPhase.PAUSED)

    # PAUSED -> AUTONOMOUS
    events.request_transition("pause_resume")
    old, new = events.consume_transition()
    assert (old, new) == (DAggerPhase.PAUSED, DAggerPhase.AUTONOMOUS)


def test_dagger_invalid_transition_ignored():
    from lerobot.rollout.strategies import DAggerEvents, DAggerPhase

    events = DAggerEvents()
    events.request_transition("correction")  # Not valid from AUTONOMOUS
    assert events.consume_transition() is None
    assert events.phase == DAggerPhase.AUTONOMOUS


def test_dagger_events_reset():
    from lerobot.rollout.strategies import DAggerEvents, DAggerPhase

    events = DAggerEvents()
    events.request_transition("pause_resume")
    events.consume_transition()  # -> PAUSED
    events.upload_requested.set()
    events.reset()
    assert events.phase == DAggerPhase.AUTONOMOUS
    assert not events.upload_requested.is_set()


# ---------------------------------------------------------------------------
# Context dataclass
# ---------------------------------------------------------------------------


def test_rollout_context_fields():
    from lerobot.rollout import RolloutContext

    field_names = {f.name for f in dataclasses.fields(RolloutContext)}
    assert field_names == {"runtime", "hardware", "policy", "processors", "data"}