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


# ---------------------------------------------------------------------------
# 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"}


# ---------------------------------------------------------------------------
# Sync engine: relative-action anchoring (drift-free chunk execution)
# ---------------------------------------------------------------------------

_REL_ACTION_NAMES = ["j0.pos", "j1.pos", "j2.pos", "gripper.pos"]
_REL_ACTION_DIM = len(_REL_ACTION_NAMES)


def _relative_pre_post(exclude_joints=None):
    """Pre/post processors wrapping the real relative (caches anchor) and absolute
    (relative + cached state) steps, mirroring what the sync engine feeds them."""
    from lerobot.processor import (
        AbsoluteActionsProcessorStep,
        RelativeActionsProcessorStep,
        TransitionKey,
        create_transition,
    )
    from lerobot.utils.constants import OBS_STATE

    relative_step = RelativeActionsProcessorStep(
        enabled=True, exclude_joints=exclude_joints or [], action_names=list(_REL_ACTION_NAMES)
    )
    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)

    class _Pre:
        steps = [relative_step]

        def __call__(self, observation):
            # Run the relative step so it caches the anchor, then pass the batch through.
            transition = create_transition(observation={OBS_STATE: observation[OBS_STATE]})
            relative_step(transition)
            return observation

        def reset(self):
            pass

    class _Post:
        def __call__(self, action):
            transition = create_transition(action=action)
            return absolute_step(transition)[TransitionKey.ACTION]

        def reset(self):
            pass

    return _Pre(), _Post(), relative_step


def _fake_relative_policy(chunk_rel, n_action_steps, chunking=True):
    """Fake relative-action policy for the sync engine.

    ``chunking=True`` buffers a chunk and serves it one action per tick, calling the
    public ``predict_action_chunk`` only on refill (pi0/fastwam/lingbot). ``False``
    returns an action directly and never calls it. The engine's anchor probe keys off
    that public call, so the fake routes through it rather than any private queue.
    """
    from collections import deque

    policy = MagicMock()
    policy.config.use_amp = False
    policy.config.action_feature_names = list(_REL_ACTION_NAMES)
    state = {"predict_calls": 0}
    queue = deque(maxlen=n_action_steps)

    def predict_action_chunk(_batch=None, **_kwargs):
        state["predict_calls"] += 1
        return chunk_rel.unsqueeze(0)  # [B=1, n, dim]

    def select_action(_observation):
        if not chunking:
            return chunk_rel[0].unsqueeze(0)
        if len(queue) == 0:
            actions = policy.predict_action_chunk(_observation)
            queue.extend(actions.transpose(0, 1))  # [n, 1, dim]
        return queue.popleft()

    policy.predict_action_chunk.side_effect = predict_action_chunk
    policy.select_action.side_effect = select_action
    policy.reset.side_effect = queue.clear
    policy._predict_state = state
    return policy


def _build_sync_engine(policy, pre, post):
    from lerobot.rollout import SyncInferenceEngine

    return SyncInferenceEngine(
        policy=policy,
        preprocessor=pre,
        postprocessor=post,
        dataset_features={"action": {"names": list(_REL_ACTION_NAMES)}},
        ordered_action_keys=list(_REL_ACTION_NAMES),
        task="test",
        device="cpu",
        robot_type="mock",
    )


def _obs_frame(state_values):
    import numpy as np

    return {"observation.state": np.asarray(state_values, dtype=np.float32)}


def test_sync_relative_holds_anchor_across_chunk():
    """Every action popped within a chunk must anchor to the tick-0 state (no drift)."""
    n = 4
    # A distinct relative offset per chunk step so a wrong anchor would be visible.
    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i in range(n)])
    pre, post, relative_step = _relative_pre_post()
    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
    engine = _build_sync_engine(policy, pre, post)

    assert engine._relative_step is relative_step  # introspection wired the step

    s0 = [1.0, 2.0, 3.0, 4.0]
    outputs = []
    for tick in range(n):
        # Feed a *different* state each tick; a drifting anchor would use it.
        state = [v + tick for v in s0]
        outputs.append(engine.get_action(_obs_frame(state)))

    # Exactly one chunk was predicted across the n ticks.
    assert policy._predict_state["predict_calls"] == 1
    for tick in range(n):
        expected = torch.tensor(s0) + chunk_rel[tick]
        torch.testing.assert_close(outputs[tick], expected)

    # Next tick empties the queue -> fresh chunk -> anchor advances to the new state.
    s_next = [10.0, 20.0, 30.0, 40.0]
    out = engine.get_action(_obs_frame(s_next))
    assert policy._predict_state["predict_calls"] == 2
    torch.testing.assert_close(out, torch.tensor(s_next) + chunk_rel[0])
    # The anchor now reflects the fresh-chunk state, not the held one.
    torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_next]))


def test_sync_relative_reset_reanchors_new_episode():
    """After ``reset()`` the first tick of the next episode anchors to the new state."""
    n = 3
    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
    pre, post, relative_step = _relative_pre_post()
    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
    engine = _build_sync_engine(policy, pre, post)

    # Episode 1: one tick anchors to s0 and leaves cached actions in the queue.
    engine.get_action(_obs_frame([1.0, 1.0, 1.0, 1.0]))
    assert policy._predict_state["predict_calls"] == 1

    engine.reset()  # clears the queue and the per-episode chunk flags

    # Episode 2: a fresh state must produce a fresh chunk anchored to that state,
    # not carry over the previous episode's anchor.
    s_new = [7.0, 8.0, 9.0, 10.0]
    out = engine.get_action(_obs_frame(s_new))
    assert policy._predict_state["predict_calls"] == 2
    torch.testing.assert_close(out, torch.tensor(s_new) + chunk_rel[0])
    torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_new]))


def test_sync_relative_non_chunking_policy_refreshes_every_tick():
    """A policy that never calls ``predict_action_chunk`` must not freeze the anchor."""
    n = 3
    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.5) for _ in range(n)])
    pre, post, _ = _relative_pre_post()
    policy = _fake_relative_policy(chunk_rel, n_action_steps=n, chunking=False)
    engine = _build_sync_engine(policy, pre, post)

    s0 = [1.0, 1.0, 1.0, 1.0]
    for tick in range(3):
        state = [v + tick for v in s0]
        out = engine.get_action(_obs_frame(state))
        # Anchor must track the current state every tick (no chunk => no hold).
        torch.testing.assert_close(out, torch.tensor(state) + chunk_rel[0])
    assert policy._predict_state["predict_calls"] == 0


def test_sync_engine_no_relative_step_is_none():
    """Without an enabled relative step, the engine takes the plain select_action path."""
    policy = MagicMock()
    policy.config.use_amp = False
    engine = _build_sync_engine(policy, MagicMock(steps=[]), MagicMock())
    assert engine._relative_step is None


def test_sync_relative_stop_restores_policy_method():
    """``stop()`` un-patches the probe so the policy object isn't permanently modified."""
    n = 3
    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
    pre, post, _ = _relative_pre_post()
    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
    original = policy.predict_action_chunk
    engine = _build_sync_engine(policy, pre, post)
    assert policy.predict_action_chunk is not original  # probe installed
    engine.stop()
    assert policy.predict_action_chunk is original  # restored