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

"""Tests for interactive rollout control: the programmatic RolloutController
and the stdin-driven InteractiveSession (--interactive=true)."""

from __future__ import annotations

import contextlib
import io
import logging
import os
import time
from threading import Event, Thread
from types import SimpleNamespace
from unittest.mock import MagicMock

import numpy as np
import pytest

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

from lerobot.rollout import (  # noqa: E402
    InferenceEngine,
    InteractiveCommand,
    InteractiveSession,
    LinkedEvent,
    RolloutController,
    RolloutEvent,
    parse_command,
)


def _wait_for(predicate, timeout: float = 2.0) -> bool:
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(0.005)
    return predicate()


@contextlib.contextmanager
def _pipe_stream():
    """A held-open pipe so the session's stdin listener never sees EOF."""
    read_fd, write_fd = os.pipe()
    reader = os.fdopen(read_fd, "r")
    writer = os.fdopen(write_fd, "w")
    try:
        yield reader, writer
    finally:
        with contextlib.suppress(OSError, ValueError):
            writer.close()
        with contextlib.suppress(OSError, ValueError):
            reader.close()


# ---------------------------------------------------------------------------
# Command parser
# ---------------------------------------------------------------------------


def test_parse_command_basic():
    cmd = parse_command("/start")
    assert cmd is not None
    assert cmd.name == "start"
    assert cmd.args == ""


def test_parse_command_case_whitespace_and_args():
    cmd = parse_command("  /SubTask Grab the red cube  ")
    assert cmd is not None
    assert cmd.name == "subtask"
    assert cmd.args == "Grab the red cube"


def test_parse_command_tab_separated():
    assert parse_command("/subtask\tgrab the cube") == InteractiveCommand(
        name="subtask", args="grab the cube"
    )


def test_parse_command_non_commands():
    assert parse_command("hello robot") is None
    assert parse_command("") is None
    assert parse_command("   ") is None
    assert parse_command("/") is None
    assert parse_command("/ start") is None


# ---------------------------------------------------------------------------
# LinkedEvent
# ---------------------------------------------------------------------------


def test_linked_event_local_flag():
    parent = Event()
    event = LinkedEvent(parent)
    assert not event.is_set()

    event.set()
    assert event.is_set()
    assert not parent.is_set()

    event.clear()
    assert not event.is_set()


def test_linked_event_reflects_parent():
    parent = Event()
    event = LinkedEvent(parent)
    parent.set()
    assert event.is_set()
    # Clearing the local flag never masks the parent.
    event.clear()
    assert event.is_set()


def test_linked_event_wait():
    parent = Event()
    event = LinkedEvent(parent)
    assert event.wait(timeout=0.05) is False

    parent.set()
    assert event.wait(timeout=0.05) is True

    parent.clear()
    event.set()
    assert event.wait(timeout=0.05) is True


def test_linked_event_wait_wakes_on_parent_set():
    parent = Event()
    event = LinkedEvent(parent)
    Thread(target=lambda: (time.sleep(0.05), parent.set()), daemon=True).start()
    assert event.wait(timeout=2.0) is True


# ---------------------------------------------------------------------------
# Shared fakes
# ---------------------------------------------------------------------------


class _FakeEngine(InferenceEngine):
    """Real task-holder semantics with mocked lifecycle methods.

    Subclassing the ABC (instead of using a bare MagicMock) means these
    tests exercise the actual ``set_task``/``task`` plumbing.
    ``failed``/``failure_traceback`` shadow the base properties as plain
    class attributes so tests can assign them.
    """

    failed = False
    failure_traceback = None

    # Declared here to satisfy the ABC; the instances below shadow them.
    def start(self) -> None: ...
    def stop(self) -> None: ...
    def reset(self) -> None: ...
    def get_action(self, obs_frame=None): ...

    def __init__(self, task: str = "pick up the cube") -> None:
        super().__init__(task=task)
        self.start = MagicMock()
        self.stop = MagicMock()
        self.reset = MagicMock()
        self.pause = MagicMock()
        self.resume = MagicMock()
        self.notify_observation = MagicMock()
        self.get_action = MagicMock(return_value=None)


def _make_ctx(run_behavior=None):
    """A mock strategy plus the minimal fake context the controller needs."""
    parent = Event()
    stop_event = LinkedEvent(parent)
    engine = _FakeEngine()
    ctx = SimpleNamespace(
        runtime=SimpleNamespace(
            cfg=SimpleNamespace(play_sounds=False),
            shutdown_event=stop_event,
        ),
        policy=SimpleNamespace(inference=engine),
        hardware=SimpleNamespace(initial_position={"joint.pos": 0.0}),
    )

    strategy = MagicMock()
    run_started = Event()

    def default_run(c):
        run_started.set()
        while not c.runtime.shutdown_event.is_set():
            time.sleep(0.005)

    strategy.run.side_effect = run_behavior or default_run
    return ctx, strategy, engine, parent, run_started


# ---------------------------------------------------------------------------
# RolloutController (the programmatic API)
# ---------------------------------------------------------------------------


def _make_controller(run_behavior=None):
    ctx, strategy, engine, parent, run_started = _make_ctx(run_behavior)
    events: list[RolloutEvent] = []
    controller = RolloutController(strategy, ctx, on_event=events.append)
    return controller, events, strategy, engine, parent, run_started


def _serve_thread(controller) -> Thread:
    thread = Thread(target=controller.serve, daemon=True)
    thread.start()
    return thread


def test_controller_requires_linked_event():
    ctx = SimpleNamespace(runtime=SimpleNamespace(shutdown_event=Event()))
    with pytest.raises(TypeError, match="LinkedEvent"):
        RolloutController(MagicMock(), ctx)


def test_controller_start_reset_stop_flow_and_events():
    controller, events, strategy, engine, _parent, run_started = _make_controller()
    thread = _serve_thread(controller)

    # Idle until start(): the strategy loop must not run on its own.
    time.sleep(0.05)
    strategy.run.assert_not_called()
    assert not controller.running

    assert controller.start() is True
    assert _wait_for(run_started.is_set)
    assert controller.running
    assert strategy.reset_control_state.call_count == 1
    assert RolloutEvent.SEGMENT_STARTED in events

    # reset() ends the segment, pauses the engine, returns the robot home.
    controller.reset()
    assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
    assert engine.pause.call_count >= 1
    assert thread.is_alive()
    assert _wait_for(lambda: RolloutEvent.RESET_DONE in events)
    assert RolloutEvent.RESET_STARTED in events

    # start() again runs a fresh segment with freshly reset control state.
    run_started.clear()
    assert controller.start() is True
    assert _wait_for(run_started.is_set)
    assert strategy.run.call_count == 2
    assert strategy.reset_control_state.call_count == 2

    # stop() ends serve(); teardown stays with the caller.
    controller.stop()
    thread.join(timeout=2.0)
    assert not thread.is_alive()
    strategy.teardown.assert_not_called()
    assert events[-1] is RolloutEvent.STOPPED


def test_controller_start_rejected_during_segment_startup():
    """A start() racing the segment startup must be rejected, not queued.

    Otherwise the re-armed request survives the whole segment and the robot
    would start again, uncommanded, when the segment ends on its own.
    """
    ctx, strategy, _engine, _parent, run_started = _make_ctx()
    in_startup = Event()
    startup_gate = Event()

    def slow_reset_control_state():
        in_startup.set()
        startup_gate.wait(timeout=2.0)

    strategy.reset_control_state.side_effect = slow_reset_control_state
    controller = RolloutController(strategy, ctx)
    thread = _serve_thread(controller)

    assert controller.start() is True
    assert _wait_for(in_startup.is_set)
    # The serve thread is inside reset_control_state: the segment counts as
    # running for concurrent callers even though strategy.run hasn't begun.
    assert controller.start() is False
    startup_gate.set()
    assert _wait_for(run_started.is_set)

    # Ending the segment must not trigger a phantom second segment.
    controller.reset()
    assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
    time.sleep(0.05)
    assert strategy.run.call_count == 1

    controller.stop()
    thread.join(timeout=2.0)


def test_controller_start_returns_false_while_running():
    controller, _events, strategy, _engine, _parent, run_started = _make_controller()
    thread = _serve_thread(controller)

    assert controller.start() is True
    assert _wait_for(run_started.is_set)
    assert controller.start() is False
    time.sleep(0.05)
    assert strategy.run.call_count == 1

    controller.stop()
    thread.join(timeout=2.0)


def test_controller_set_task_and_reset_restores_launch_task():
    controller, _events, strategy, engine, _parent, _run_started = _make_controller()
    thread = _serve_thread(controller)
    initial = controller.initial_task

    # reset() with the launch task still in place reports no restore.
    assert controller.reset() is False
    assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)

    assert controller.set_task("fold the towel") is True
    assert controller.task == "fold the towel"
    assert engine.task == "fold the towel"
    assert controller.set_task("fold the towel") is False

    assert controller.reset() is True  # the task had been changed
    assert controller.task == initial

    controller.stop()
    thread.join(timeout=2.0)


def test_controller_engine_failure_emits_event():
    def failing_run(c):
        c.policy.inference.failed = True
        c.policy.inference.failure_traceback = "RuntimeError: boom-traceback"
        c.runtime.shutdown_event.set()

    controller, events, strategy, _engine, _parent, _run_started = _make_controller(failing_run)
    thread = _serve_thread(controller)
    controller.start()
    thread.join(timeout=2.0)
    assert not thread.is_alive()
    strategy.return_to_initial_position.assert_not_called()
    assert RolloutEvent.ENGINE_FAILED in events
    assert controller.failed
    assert controller.failure_traceback == "RuntimeError: boom-traceback"


def test_controller_segment_ended_event_on_natural_end():
    def finite_run(c):
        return None  # e.g. --duration elapsed

    controller, events, strategy, _engine, _parent, _run_started = _make_controller(finite_run)
    thread = _serve_thread(controller)

    controller.start()
    assert _wait_for(lambda: RolloutEvent.SEGMENT_ENDED in events)
    assert thread.is_alive()  # back to idle, not shut down
    assert not controller.running

    controller.stop()
    thread.join(timeout=2.0)


def test_controller_reset_skipped_without_initial_position():
    ctx, strategy, _engine, _parent, _run_started = _make_ctx()
    ctx.hardware.initial_position = {}
    events: list[RolloutEvent] = []
    controller = RolloutController(strategy, ctx, on_event=events.append)
    thread = _serve_thread(controller)

    controller.reset()
    assert _wait_for(lambda: RolloutEvent.RESET_SKIPPED in events)
    strategy.return_to_initial_position.assert_not_called()

    controller.stop()
    thread.join(timeout=2.0)


def test_controller_callback_errors_do_not_kill_serve():
    ctx, strategy, _engine, _parent, run_started = _make_ctx()

    def broken_observer(event):
        raise RuntimeError("observer boom")

    controller = RolloutController(strategy, ctx, on_event=broken_observer)
    thread = _serve_thread(controller)

    controller.start()
    assert _wait_for(run_started.is_set)
    controller.stop()
    thread.join(timeout=2.0)
    assert not thread.is_alive()


def test_controller_works_without_event_callback():
    ctx, strategy, _engine, _parent, _run_started = _make_ctx()
    controller = RolloutController(strategy, ctx)
    thread = _serve_thread(controller)
    controller.start()
    controller.stop()
    thread.join(timeout=2.0)
    assert not thread.is_alive()


# ---------------------------------------------------------------------------
# InteractiveSession (the stdin CLI front-end)
# ---------------------------------------------------------------------------


def _make_session(input_stream, run_behavior=None):
    """Build a session around a mock strategy and a minimal fake context."""
    ctx, strategy, engine, parent, run_started = _make_ctx(run_behavior)
    session = InteractiveSession(strategy, ctx, input_stream=input_stream)
    return session, strategy, engine, parent, run_started


def _start_session_thread(session) -> Thread:
    thread = Thread(target=session.run, daemon=True)
    thread.start()
    return thread


def test_session_requires_linked_event():
    ctx = SimpleNamespace(runtime=SimpleNamespace(shutdown_event=Event()))
    with pytest.raises(TypeError, match="LinkedEvent"):
        InteractiveSession(MagicMock(), ctx)


def test_session_start_reset_restart_stop_flow():
    with _pipe_stream() as (reader, _writer):
        session, strategy, engine, _parent, run_started = _make_session(reader)
        thread = _start_session_thread(session)

        # Idle until /start: the strategy loop must not run on its own.
        time.sleep(0.05)
        strategy.run.assert_not_called()

        session._handle_line("/start")
        assert _wait_for(run_started.is_set)
        assert strategy.reset_control_state.call_count == 1

        # /reset ends the segment, pauses the engine, and returns to the initial position.
        session._handle_line("/reset")
        assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
        assert engine.pause.call_count >= 1
        assert thread.is_alive()

        # /start again runs a fresh segment with freshly reset control state.
        run_started.clear()
        session._handle_line("/start")
        assert _wait_for(run_started.is_set)
        assert strategy.run.call_count == 2
        assert strategy.reset_control_state.call_count == 2

        # /stop ends the session; teardown stays with the caller (the CLI script).
        session._handle_line("/stop")
        thread.join(timeout=2.0)
        assert not thread.is_alive()
        strategy.teardown.assert_not_called()


def test_session_stop_while_idle():
    with _pipe_stream() as (reader, _writer):
        session, strategy, _engine, _parent, _run_started = _make_session(reader)
        thread = _start_session_thread(session)
        session._handle_line("/stop")
        thread.join(timeout=2.0)
        assert not thread.is_alive()
        strategy.run.assert_not_called()


def test_session_reset_while_idle_returns_to_initial_position():
    with _pipe_stream() as (reader, _writer):
        session, strategy, _engine, _parent, _run_started = _make_session(reader)
        thread = _start_session_thread(session)
        session._handle_line("/reset")
        assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
        assert thread.is_alive()
        session._handle_line("/stop")
        thread.join(timeout=2.0)


def test_session_start_while_running_is_rejected():
    with _pipe_stream() as (reader, _writer):
        session, strategy, _engine, _parent, run_started = _make_session(reader)
        thread = _start_session_thread(session)

        session._handle_line("/start")
        assert _wait_for(run_started.is_set)
        session._handle_line("/start")
        time.sleep(0.05)
        assert strategy.run.call_count == 1

        session._handle_line("/stop")
        thread.join(timeout=2.0)


def test_session_reset_cancels_pending_start():
    """Last command wins: a queued /start must not fire after a later /reset."""
    with _pipe_stream() as (reader, _writer):
        session, strategy, _engine, _parent, _run_started = _make_session(reader)
        # Queue both commands before the session loop starts servicing them.
        session._handle_line("/start")
        session._handle_line("/reset")
        thread = _start_session_thread(session)

        assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
        time.sleep(0.05)
        strategy.run.assert_not_called()

        session._handle_line("/stop")
        thread.join(timeout=2.0)


def test_session_stop_cancels_pending_start():
    with _pipe_stream() as (reader, _writer):
        session, strategy, _engine, _parent, _run_started = _make_session(reader)
        session._handle_line("/start")
        session._handle_line("/stop")
        thread = _start_session_thread(session)
        thread.join(timeout=2.0)
        assert not thread.is_alive()
        strategy.run.assert_not_called()


def test_session_exits_on_parent_shutdown():
    with _pipe_stream() as (reader, _writer):
        session, _strategy, _engine, parent, run_started = _make_session(reader)
        thread = _start_session_thread(session)
        session._handle_line("/start")
        assert _wait_for(run_started.is_set)

        parent.set()  # SIGINT/SIGTERM path
        thread.join(timeout=2.0)
        assert not thread.is_alive()


def test_session_stops_on_engine_failure(capsys):
    def failing_run(c):
        # Mimic the RTC thread's fatal-error path: flag the failure, capture
        # the traceback, and set the engine's shutdown event (the LinkedEvent).
        c.policy.inference.failed = True
        c.policy.inference.failure_traceback = "RuntimeError: boom-traceback"
        c.runtime.shutdown_event.set()

    with _pipe_stream() as (reader, _writer):
        session, strategy, _engine, _parent, _run_started = _make_session(reader, run_behavior=failing_run)
        thread = _start_session_thread(session)
        session._handle_line("/start")
        thread.join(timeout=2.0)
        assert not thread.is_alive()
        # A failed engine ends the session instead of returning to idle, and
        # the captured traceback is surfaced despite the muted console logs.
        strategy.return_to_initial_position.assert_not_called()
        assert "boom-traceback" in capsys.readouterr().out


def test_session_stops_on_engine_failure_while_idle():
    """A fatal engine error while idle ends the session instead of being masked by /start."""
    with _pipe_stream() as (reader, _writer):
        session, strategy, engine, _parent, _run_started = _make_session(reader)
        thread = _start_session_thread(session)
        time.sleep(0.05)
        engine.failed = True
        thread.join(timeout=2.0)
        assert not thread.is_alive()
        strategy.run.assert_not_called()


def test_session_returns_to_idle_when_run_ends_naturally():
    def finite_run(c):
        return None  # e.g. --duration elapsed

    with _pipe_stream() as (reader, _writer):
        session, strategy, _engine, _parent, _run_started = _make_session(reader, run_behavior=finite_run)
        thread = _start_session_thread(session)

        session._handle_line("/start")
        assert _wait_for(lambda: strategy.run.call_count == 1)
        time.sleep(0.05)
        assert thread.is_alive()  # back to idle, not shut down

        # The session accepts another /start after a natural end.
        session._handle_line("/start")
        assert _wait_for(lambda: strategy.run.call_count == 2)

        session._handle_line("/stop")
        thread.join(timeout=2.0)


def test_session_unknown_input_does_not_start(capsys):
    with _pipe_stream() as (reader, _writer):
        session, strategy, _engine, _parent, _run_started = _make_session(reader)
        thread = _start_session_thread(session)

        session._handle_line("/frobnicate")
        session._handle_line("hello robot")
        session._handle_line("/help")
        time.sleep(0.05)
        strategy.run.assert_not_called()

        out = capsys.readouterr().out
        assert "/frobnicate" in out
        assert "commands start with '/'" in out

        session._handle_line("/stop")
        thread.join(timeout=2.0)


def test_session_eof_stops_session():
    with _pipe_stream() as (reader, writer):
        session, strategy, _engine, _parent, _run_started = _make_session(reader)
        thread = _start_session_thread(session)
        writer.close()  # EOF on the command stream
        thread.join(timeout=2.0)
        assert not thread.is_alive()
        strategy.run.assert_not_called()


def test_session_commands_via_stream():
    """End-to-end: commands flow through the pipe and the listener thread."""
    with _pipe_stream() as (reader, writer):
        session, strategy, _engine, _parent, run_started = _make_session(reader)
        thread = _start_session_thread(session)

        writer.write("/start\n")
        writer.flush()
        assert _wait_for(run_started.is_set)

        writer.write("/stop\n")
        writer.flush()
        thread.join(timeout=2.0)
        assert not thread.is_alive()
        assert strategy.run.call_count == 1


def test_session_mutes_logs_below_error_and_restores_on_exit():
    import warnings

    # Libraries like transformers attach their own console handler with
    # propagate=False; the process-wide logging.disable gate covers those too.
    lib_stream = io.StringIO()
    lib_logger = logging.getLogger("test_interactive_fake_lib")
    lib_logger.propagate = False
    lib_logger.setLevel(logging.DEBUG)
    lib_handler = logging.StreamHandler(lib_stream)
    lib_handler.setLevel(logging.INFO)
    lib_logger.addHandler(lib_handler)
    n_warning_filters = len(warnings.filters)
    try:
        with _pipe_stream() as (reader, _writer):
            session, _strategy, _engine, _parent, _run_started = _make_session(reader)
            thread = _start_session_thread(session)
            assert _wait_for(lambda: logging.root.manager.disable == logging.WARNING)

            lib_logger.info("muted-info")
            lib_logger.warning("muted-warning")
            lib_logger.error("visible-error")  # errors must surface mid-session

            session._handle_line("/stop")
            thread.join(timeout=2.0)

        output = lib_stream.getvalue()
        assert "muted-info" not in output
        assert "muted-warning" not in output
        assert "visible-error" in output

        # Everything is restored once the session ends.
        assert logging.root.manager.disable == logging.NOTSET
        assert len(warnings.filters) == n_warning_filters
        lib_logger.info("post-session-info")
        assert "post-session-info" in lib_stream.getvalue()
    finally:
        lib_logger.removeHandler(lib_handler)


def test_session_restores_preexisting_disable_level():
    """An embedding application's own logging.disable level survives the session."""
    logging.disable(logging.DEBUG)
    try:
        with _pipe_stream() as (reader, _writer):
            session, _strategy, _engine, _parent, _run_started = _make_session(reader)
            thread = _start_session_thread(session)
            assert _wait_for(lambda: logging.root.manager.disable == logging.WARNING)
            session._handle_line("/stop")
            thread.join(timeout=2.0)
        assert logging.root.manager.disable == logging.DEBUG
    finally:
        logging.disable(logging.NOTSET)


def test_session_drives_real_base_strategy():
    """End-to-end with a real BaseStrategy control loop (only hardware/engine mocked)."""
    from lerobot.rollout import BaseStrategy, BaseStrategyConfig

    parent = Event()
    stop_event = LinkedEvent(parent)
    engine = _FakeEngine()
    engine.get_action.return_value = None  # no action ready; the loop still ticks

    robot = MagicMock()
    robot.get_observation.return_value = {"joint.pos": 0.0}

    def identity(x):
        return x

    ctx = SimpleNamespace(
        runtime=SimpleNamespace(
            cfg=SimpleNamespace(
                play_sounds=False,
                fps=100.0,
                duration=0.0,
                use_torch_compile=False,
                interpolation_multiplier=1,
                display_data=False,
            ),
            shutdown_event=stop_event,
        ),
        policy=SimpleNamespace(inference=engine),
        hardware=SimpleNamespace(robot_wrapper=robot, teleop=None, initial_position={"joint.pos": 0.0}),
        processors=SimpleNamespace(
            teleop_action_processor=identity,
            robot_action_processor=identity,
            robot_observation_processor=identity,
        ),
        data=SimpleNamespace(dataset=None, dataset_features={}, hw_features={}, ordered_action_keys=[]),
    )

    strategy = BaseStrategy(BaseStrategyConfig())
    strategy.setup(ctx)
    strategy.return_to_initial_position = MagicMock()  # skip the 3s hardware sweep

    with _pipe_stream() as (reader, _writer):
        session = InteractiveSession(strategy, ctx, input_stream=reader)
        thread = _start_session_thread(session)

        session._handle_line("/start")
        assert _wait_for(lambda: engine.resume.called)
        assert _wait_for(lambda: robot.get_observation.call_count >= 3)

        session._handle_line("/reset")
        assert _wait_for(lambda: strategy.return_to_initial_position.called)
        assert engine.pause.called
        assert thread.is_alive()

        session._handle_line("/start")
        assert _wait_for(lambda: engine.resume.call_count >= 2)

        session._handle_line("/stop")
        thread.join(timeout=2.0)
        assert not thread.is_alive()


def test_session_subtask_sets_and_reports_task(capsys):
    with _pipe_stream() as (reader, _writer):
        session, _strategy, engine, _parent, _run_started = _make_session(reader)
        thread = _start_session_thread(session)

        session._handle_line("/subtask grab the red cube")
        assert engine.task == "grab the red cube"

        # No argument reports the current task without changing it.
        session._handle_line("/subtask")
        assert engine.task == "grab the red cube"

        # Re-issuing the same task is reported as a no-op.
        session._handle_line("/subtask grab the red cube")
        out = capsys.readouterr().out
        assert "Current task: 'grab the red cube'" in out
        assert "Task unchanged: 'grab the red cube'" in out

        session._handle_line("/stop")
        thread.join(timeout=2.0)


def test_session_subtask_strips_quotes_and_works_while_running():
    with _pipe_stream() as (reader, _writer):
        session, _strategy, engine, _parent, run_started = _make_session(reader)
        thread = _start_session_thread(session)

        session._handle_line("/start")
        assert _wait_for(run_started.is_set)

        # Switching mid-run must not interrupt the control loop.
        session._handle_line('/subtask "fold the towel"')
        assert engine.task == "fold the towel"
        time.sleep(0.05)
        assert session.controller.running

        session._handle_line("/stop")
        thread.join(timeout=2.0)


def test_session_subtask_after_reset_is_not_clobbered():
    """A /subtask issued right after /reset must win — both writes are ordered by command order."""
    with _pipe_stream() as (reader, writer):
        session, strategy, engine, _parent, run_started = _make_session(reader)
        thread = _start_session_thread(session)

        writer.write("/start\n")
        writer.flush()
        assert _wait_for(run_started.is_set)

        # Arrives as one chunk, so both handlers run back-to-back on the
        # listener thread while the segment is still unwinding.
        writer.write("/reset\n/subtask put the cube in the box\n")
        writer.flush()
        assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
        time.sleep(0.1)
        assert engine.task == "put the cube in the box"

        session._handle_line("/stop")
        thread.join(timeout=2.0)


def test_session_reset_restores_initial_task():
    with _pipe_stream() as (reader, _writer):
        session, strategy, engine, _parent, _run_started = _make_session(reader)
        initial = engine.task
        thread = _start_session_thread(session)

        session._handle_line("/subtask fold the towel")
        assert engine.task == "fold the towel"

        session._handle_line("/reset")
        assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
        assert engine.task == initial

        session._handle_line("/stop")
        thread.join(timeout=2.0)


# ---------------------------------------------------------------------------
# InferenceEngine task holder (the /subtask plumbing)
# ---------------------------------------------------------------------------


def test_engine_task_holder_tracks_changes():
    engine = _FakeEngine("pick up the cube")
    assert engine.task == "pick up the cube"
    # No change yet: nothing for the inference thread to flush.
    assert engine._take_task() == ("pick up the cube", False)

    assert engine.set_task("fold the towel") is True
    assert engine.task == "fold the towel"
    # The change edge is delivered once, then consumed.
    assert engine._take_task() == ("fold the towel", True)
    assert engine._take_task() == ("fold the towel", False)

    # Setting the same value is a no-op and raises no edge.
    assert engine.set_task("fold the towel") is False
    assert engine._take_task() == ("fold the towel", False)


def test_engine_discard_task_change():
    engine = _FakeEngine("a")
    engine.set_task("b")
    engine._discard_task_change()
    assert engine._take_task() == ("b", False)


def test_sync_engine_uses_new_task_and_flushes_precomputed_actions():
    """A /subtask switch must reach the policy and drop stale queued actions."""
    import torch

    from lerobot.rollout.inference import SyncInferenceEngine

    policy = MagicMock()
    policy.config.use_amp = False
    policy.select_action.return_value = torch.zeros(1, 2)

    engine = SyncInferenceEngine(
        policy=policy,
        preprocessor=lambda obs: obs,
        postprocessor=lambda action: action,
        dataset_features={
            "action": {"dtype": "float32", "shape": (2,), "names": ["j1.pos", "j2.pos"]},
        },
        ordered_action_keys=["j1.pos", "j2.pos"],
        task="pick up the cube",
        device="cpu",
        robot_type="mock",
    )

    engine.get_action({"observation.state": np.zeros(1, dtype=np.float32)})
    assert policy.drop_queued_actions.call_count == 0
    assert policy.select_action.call_args[0][0]["task"] == "pick up the cube"

    engine.set_task("fold the towel")
    engine.get_action({"observation.state": np.zeros(1, dtype=np.float32)})
    # Precomputed chunk actions are dropped so the new task applies immediately,
    # without the wider episode reset (which would perturb observation history).
    assert policy.drop_queued_actions.call_count == 1
    assert policy.reset.call_count == 0
    assert policy.select_action.call_args[0][0]["task"] == "fold the towel"

    # Only the first call after a switch flushes.
    engine.get_action({"observation.state": np.zeros(1, dtype=np.float32)})
    assert policy.drop_queued_actions.call_count == 1


def test_drop_queued_actions_clears_both_queue_conventions():
    """PreTrainedPolicy.drop_queued_actions covers both action-queue idioms in the repo."""
    from collections import deque

    from lerobot.policies.pretrained import PreTrainedPolicy
    from lerobot.utils.constants import ACTION

    # PreTrainedPolicy's metaclass demands a config_class, so exercise the
    # method against stand-ins carrying each queue idiom.
    flush = PreTrainedPolicy.drop_queued_actions

    queues_policy = SimpleNamespace(  # smolvla / diffusion / vqbet / wall_x style
        _queues={ACTION: deque([1, 2, 3]), "observation.state": deque([9])}
    )
    flush(queues_policy)
    assert len(queues_policy._queues[ACTION]) == 0
    # Other episode state is intentionally left alone.
    assert len(queues_policy._queues["observation.state"]) == 1

    action_queue_policy = SimpleNamespace(_action_queue=deque([1, 2, 3]))  # act / pi0 / groot style
    flush(action_queue_policy)
    assert len(action_queue_policy._action_queue) == 0

    # Queue-less policies inherit a no-op.
    flush(SimpleNamespace())


# ---------------------------------------------------------------------------
# Sentry strategy: restartable run() segments + live task labels
# ---------------------------------------------------------------------------


def _make_sentry(monkeypatch):
    from lerobot.rollout import SentryStrategy, SentryStrategyConfig

    # Keep setup independent of camera features and never rotate episodes
    # mid-test; frame assembly is not under test here.
    monkeypatch.setattr("lerobot.rollout.strategies.sentry.estimate_max_episode_seconds", lambda *a, **k: 1e9)
    monkeypatch.setattr(
        "lerobot.rollout.strategies.sentry.send_next_action", lambda *a, **k: {"joint.pos": 0.5}
    )
    monkeypatch.setattr("lerobot.rollout.strategies.sentry.build_dataset_frame", lambda *a, **k: {})

    engine = _FakeEngine()
    dataset = MagicMock()
    robot = MagicMock()
    robot.get_observation.return_value = {"joint.pos": 0.0}
    stop_event = Event()

    def identity(x):
        return x

    ctx = SimpleNamespace(
        runtime=SimpleNamespace(
            cfg=SimpleNamespace(
                play_sounds=False,
                fps=200.0,
                duration=0.0,
                use_torch_compile=False,
                interpolation_multiplier=1,
                display_data=False,
                return_to_initial_position=False,
                task="pick up the cube",
                dataset=SimpleNamespace(push_to_hub=False, tags=None, private=False),
            ),
            shutdown_event=stop_event,
        ),
        policy=SimpleNamespace(inference=engine),
        hardware=SimpleNamespace(robot_wrapper=robot, teleop=None, initial_position={"joint.pos": 0.0}),
        processors=SimpleNamespace(
            teleop_action_processor=identity,
            robot_action_processor=identity,
            robot_observation_processor=identity,
        ),
        data=SimpleNamespace(dataset=dataset, dataset_features={}, hw_features={}, ordered_action_keys=[]),
    )

    strategy = SentryStrategy(SentryStrategyConfig())
    strategy.setup(ctx)
    return strategy, ctx, dataset, engine, stop_event


def test_sentry_run_is_restartable_and_finalizes_only_in_teardown(monkeypatch):
    strategy, ctx, dataset, _engine, stop_event = _make_sentry(monkeypatch)

    # Segment 1.
    thread = Thread(target=strategy.run, args=(ctx,), daemon=True)
    thread.start()
    assert _wait_for(lambda: dataset.add_frame.call_count >= 2)
    stop_event.set()
    thread.join(timeout=2.0)
    assert not thread.is_alive()

    # The segment saved its partial episode but left the dataset open.
    assert dataset.save_episode.call_count == 1
    dataset.finalize.assert_not_called()

    # Segment 2: run() is restartable on the same instance.
    stop_event.clear()
    frames_before = dataset.add_frame.call_count
    thread = Thread(target=strategy.run, args=(ctx,), daemon=True)
    thread.start()
    assert _wait_for(lambda: dataset.add_frame.call_count >= frames_before + 2)
    stop_event.set()
    thread.join(timeout=2.0)
    assert dataset.save_episode.call_count == 2
    dataset.finalize.assert_not_called()

    # Only teardown finalizes.
    strategy.teardown(ctx)
    dataset.finalize.assert_called_once()


def test_sentry_labels_frames_with_live_engine_task(monkeypatch):
    strategy, ctx, dataset, engine, stop_event = _make_sentry(monkeypatch)

    thread = Thread(target=strategy.run, args=(ctx,), daemon=True)
    thread.start()
    assert _wait_for(lambda: dataset.add_frame.call_count >= 2)

    # A live task change (e.g. /subtask through the controller) relabels
    # recorded frames from that moment on.
    engine.set_task("fold the towel")
    assert _wait_for(
        lambda: any(call.args[0]["task"] == "fold the towel" for call in dataset.add_frame.call_args_list)
    )
    stop_event.set()
    thread.join(timeout=2.0)

    tasks = [call.args[0]["task"] for call in dataset.add_frame.call_args_list]
    assert tasks[0] == "pick up the cube"
    assert tasks[-1] == "fold the towel"

    strategy.teardown(ctx)


def test_sentry_discards_tail_episode_when_final_save_fails(monkeypatch):
    strategy, ctx, dataset, _engine, stop_event = _make_sentry(monkeypatch)
    dataset.save_episode.side_effect = ValueError("disk full mid-write")

    thread = Thread(target=strategy.run, args=(ctx,), daemon=True)
    thread.start()
    assert _wait_for(lambda: dataset.add_frame.call_count >= 1)
    stop_event.set()
    thread.join(timeout=2.0)
    assert not thread.is_alive()

    # The failed tail save must not leave a half-written streaming encode for
    # teardown's finalize to flush, nor a half-mutated episode buffer that
    # would crash the first add_frame of a restarted segment.
    dataset.writer.cancel_pending_videos.assert_called_once()
    assert dataset.writer.episode_buffer is None
    dataset.finalize.assert_not_called()


# ---------------------------------------------------------------------------
# Config validation
# ---------------------------------------------------------------------------


def test_interactive_rejects_keyboard_bound_strategies():
    from lerobot.configs.dataset import DatasetRecordConfig
    from lerobot.rollout import HighlightStrategyConfig, RolloutConfig
    from tests.mocks.mock_robot import MockRobotConfig

    with pytest.raises(ValueError, match="--interactive=true supports"):
        RolloutConfig(
            robot=MockRobotConfig(),
            strategy=HighlightStrategyConfig(),
            dataset=DatasetRecordConfig(repo_id="user/rollout_test", single_task="test"),
            interactive=True,
        )


def test_interactive_allows_sentry():
    from lerobot.configs.dataset import DatasetRecordConfig
    from lerobot.rollout import RolloutConfig, SentryStrategyConfig
    from tests.mocks.mock_robot import MockRobotConfig

    cfg = RolloutConfig(
        robot=MockRobotConfig(),
        strategy=SentryStrategyConfig(),
        dataset=DatasetRecordConfig(repo_id="user/rollout_test", single_task="test"),
        policy=SimpleNamespace(device="cpu"),  # stands in for a PreTrainedConfig
        interactive=True,
    )
    assert cfg.interactive is True
    assert cfg.dataset.streaming_encoding is True  # sentry forces streaming