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

"""Programmatic control of a rollout: start, pause, re-instruct, and stop a
policy while hardware and policy stay connected and warm.

:class:`RolloutController` is the embedding-friendly core of interactive
rollouts.  It has no I/O of its own — no stdin, no printing, no log
manipulation — so it can be driven from any application code: a CLI
(:class:`lerobot.rollout.interactive.InteractiveSession` is exactly that), a
network server, a voice front-end, or a notebook.

Typical embedding::

    from threading import Event, Thread
    from lerobot.rollout import (
        LinkedEvent,
        RolloutController,
        build_rollout_context,
        create_strategy,
    )

    parent = Event()  # your application's shutdown signal
    ctx = build_rollout_context(cfg, LinkedEvent(parent))
    strategy = create_strategy(cfg.strategy)
    strategy.setup(ctx)

    controller = RolloutController(strategy, ctx)
    serve_thread = Thread(target=controller.serve)
    serve_thread.start()  # or call serve() on your main thread

    controller.start()  # robot starts executing the policy
    controller.set_task("grab the red cube")  # re-instruct mid-run
    controller.reset()  # stop movement, return home, stay warm
    controller.stop()  # end serve()

    serve_thread.join()
    strategy.teardown(ctx)  # teardown stays with the caller
"""

from __future__ import annotations

import logging
import time
from collections.abc import Callable
from enum import Enum
from threading import Event, Lock
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .context import RolloutContext
    from .strategies import RolloutStrategy

logger = logging.getLogger(__name__)


class LinkedEvent(Event):
    """A ``threading.Event`` whose ``is_set`` also reflects a parent event.

    ``set``/``clear`` act only on the local flag, so a controller can raise
    and clear its own segment-stop requests without masking (or accidentally
    re-arming) the process-wide shutdown event carried by ``parent``.  Every
    rollout strategy control loop polls ``ctx.runtime.shutdown_event.is_set()``,
    so installing a ``LinkedEvent`` there makes the loops react both to
    controller commands and to real shutdown signals.
    """

    _WAIT_SLICE_S = 0.05

    def __init__(self, parent: Event) -> None:
        super().__init__()
        self.parent = parent

    def is_set(self) -> bool:
        return super().is_set() or self.parent.is_set()

    def wait(self, timeout: float | None = None) -> bool:
        """Wait for either the local or the parent flag.

        The base ``Event.wait`` only watches the local flag, so poll in short
        slices to also observe the parent.  Strategy loops only call
        ``is_set()``; this coarse wait exists for API completeness.
        """
        deadline = None if timeout is None else time.perf_counter() + timeout
        while not self.is_set():
            remaining = None if deadline is None else deadline - time.perf_counter()
            if remaining is not None and remaining <= 0:
                return False
            wait_slice = self._WAIT_SLICE_S if remaining is None else min(self._WAIT_SLICE_S, remaining)
            super().wait(wait_slice)
        return True


class RolloutEvent(Enum):
    """Lifecycle notifications emitted by :class:`RolloutController`.

    All events are emitted on the thread running :meth:`RolloutController.serve`;
    callbacks must be quick and must not call back into the controller's
    blocking methods.
    """

    SEGMENT_STARTED = "segment_started"
    """A control-loop segment is about to run (control state freshly reset)."""

    SEGMENT_ENDED = "segment_ended"
    """The segment returned on its own (e.g. ``--duration`` elapsed); the
    controller is idle again and the robot is holding position."""

    RESET_STARTED = "reset_started"
    """A reset is being executed: inference paused, robot about to move home."""

    RESET_DONE = "reset_done"
    """The robot is back at its initial position, holding."""

    RESET_SKIPPED = "reset_skipped"
    """No initial position was captured; the robot holds its current pose."""

    ENGINE_FAILED = "engine_failed"
    """The inference engine hit an unrecoverable error; ``serve()`` is about
    to return.  Read :attr:`RolloutController.failure_traceback` for details."""

    STOPPED = "stopped"
    """``serve()`` is returning (after :meth:`RolloutController.stop`, EOF of
    the driving front-end, an engine failure, or a parent shutdown signal)."""


class RolloutController:
    """Drive a rollout strategy through thread-safe start/reset/stop/set_task calls.

    The controller owns the outer lifecycle between ``strategy.setup(ctx)``
    and ``strategy.teardown(ctx)`` (both stay with the caller): the robot is
    idle until :meth:`start`, each run *segment* executes ``strategy.run(ctx)``
    on the thread that called :meth:`serve` until interrupted or until the
    strategy returns on its own (e.g. ``--duration`` elapsed).  :meth:`reset`
    pauses the inference engine, returns the robot to its initial position,
    and restores the launch task, while hardware and policy stay warm.
    :meth:`stop` ends :meth:`serve` so the caller can run
    ``strategy.teardown(ctx)``.

    Requires ``ctx.runtime.shutdown_event`` to be a :class:`LinkedEvent`: the
    controller sets the local flag to end a segment, and process shutdown
    signals still propagate through the parent.  Build the context with
    ``build_rollout_context(cfg, LinkedEvent(shutdown_event))``.

    Thread safety: the control methods (:meth:`start`, :meth:`reset`,
    :meth:`stop`, :meth:`set_task`) may be called from any thread and are
    serialized by an internal lock, so calls issued in order from one thread
    keep that order — e.g. a ``set_task`` right after a ``reset`` is not
    clobbered by the reset's task restore.  Commands are last-write-wins:
    ``reset`` and ``stop`` cancel a still-pending ``start`` so the robot
    never starts moving after the caller's most recent command asked it not
    to.  Events are emitted on the :meth:`serve` thread via ``on_event``.
    """

    _POLL_INTERVAL_S = 0.2

    def __init__(
        self,
        strategy: RolloutStrategy,
        ctx: RolloutContext,
        on_event: Callable[[RolloutEvent], None] | None = None,
    ) -> None:
        stop_event = ctx.runtime.shutdown_event
        if not isinstance(stop_event, LinkedEvent):
            raise TypeError(
                "RolloutController requires ctx.runtime.shutdown_event to be a LinkedEvent so "
                "reset() can end a run segment without triggering process shutdown. Build the "
                "rollout context with build_rollout_context(cfg, LinkedEvent(shutdown_event))."
            )
        self._strategy = strategy
        self._ctx = ctx
        self._segment_stop = stop_event
        self._global_shutdown = stop_event.parent
        self._on_event = on_event
        # The instruction the rollout was launched with; reset() restores it.
        self._initial_task = ctx.policy.inference.task

        # Serializes the control methods so multi-writer task updates (e.g.
        # reset()'s restore followed by a set_task()) keep their call order.
        self._control_lock = Lock()

        # Written by control methods (any thread), consumed by the serve loop.
        self._start_requested = Event()
        self._reset_requested = Event()
        self._stop_requested = Event()
        self._wake = Event()
        self._running = Event()

    # ------------------------------------------------------------------
    # Introspection
    # ------------------------------------------------------------------

    @property
    def task(self) -> str:
        """The language instruction currently conditioning inference."""
        return self._ctx.policy.inference.task

    @property
    def initial_task(self) -> str:
        """The instruction the rollout was launched with (restored by :meth:`reset`)."""
        return self._initial_task

    @property
    def running(self) -> bool:
        """True while a control-loop segment is executing."""
        return self._running.is_set()

    @property
    def failed(self) -> bool:
        """True if the inference engine hit an unrecoverable error."""
        return self._ctx.policy.inference.failed

    @property
    def failure_traceback(self) -> str | None:
        """Formatted traceback of the engine failure, when :attr:`failed` is True."""
        return self._ctx.policy.inference.failure_traceback

    # ------------------------------------------------------------------
    # Control methods (callable from any thread)
    # ------------------------------------------------------------------

    def start(self) -> bool:
        """Request a control-loop segment.

        Returns ``False`` when a segment is already running (the request is
        ignored); ``True`` when the segment was scheduled.  The segment itself
        executes on the :meth:`serve` thread.
        """
        with self._control_lock:
            if self._running.is_set():
                return False
            self._start_requested.set()
            self._wake.set()
            return True

    def reset(self) -> bool:
        """Stop movement, return the robot to its initial position, restore the launch task.

        Hardware and policy stay warm; call :meth:`start` to run again.
        Returns ``True`` when the task was restored to the launch task (i.e.
        it had been changed), ``False`` when it was already the launch task.
        """
        with self._control_lock:
            # Last command wins: a start() still waiting to be serviced is
            # cancelled so the robot never starts moving after the caller
            # asked it not to.  Flag first, segment-stop second (see the
            # ordering note in _run_segment).
            self._start_requested.clear()
            # Restore the task here, under the control lock, rather than in
            # _reset_robot (which runs later, on the serve thread) so that a
            # set_task() issued right after this reset() is not silently
            # reverted by a deferred restore.
            restored = self._ctx.policy.inference.set_task(self._initial_task)
            self._reset_requested.set()
            self._segment_stop.set()
            self._wake.set()
            return restored

    def stop(self) -> None:
        """End :meth:`serve`; the caller then runs ``strategy.teardown(ctx)``."""
        with self._control_lock:
            self._start_requested.clear()  # last command wins, see reset()
            self._stop_requested.set()
            self._segment_stop.set()
            self._wake.set()

    def set_task(self, task: str) -> bool:
        """Change the instruction the policy follows, effective from the next inference.

        Returns ``True`` when the value actually changed.  Safe to call while
        a segment is running: the engine applies the switch on its own
        inference thread (sync backends also drop actions precomputed under
        the previous instruction).
        """
        with self._control_lock:
            return self._ctx.policy.inference.set_task(task)

    # ------------------------------------------------------------------
    # Serve loop (blocks the calling thread)
    # ------------------------------------------------------------------

    def serve(self) -> None:
        """Service control requests until :meth:`stop`, engine failure, or parent shutdown.

        Blocks the calling thread; run segments execute here.  Emits
        :class:`RolloutEvent` notifications through ``on_event``.
        """
        try:
            while not self._global_shutdown.is_set():
                if self._ctx.policy.inference.failed:
                    self._emit(RolloutEvent.ENGINE_FAILED)
                    break
                if self._stop_requested.is_set():
                    break
                if self._reset_requested.is_set():
                    self._reset_requested.clear()
                    self._reset_robot()
                    continue
                if self._start_requested.is_set():
                    # Consume the request and mark the segment running in one
                    # atomic step: start() gates on _running, so a concurrent
                    # start() is rejected for the entire startup sequence
                    # (reset_control_state, SEGMENT_STARTED emission), not just
                    # once strategy.run() begins — otherwise it could re-arm
                    # _start_requested behind the running segment and the robot
                    # would start again, uncommanded, when the segment ends.
                    with self._control_lock:
                        starting = self._start_requested.is_set()
                        if starting:
                            self._start_requested.clear()
                            self._running.set()
                    if starting:
                        self._run_segment()
                    continue
                self._wake.wait(timeout=self._POLL_INTERVAL_S)
                self._wake.clear()
        finally:
            self._emit(RolloutEvent.STOPPED)

    def _run_segment(self) -> None:
        """Execute one ``strategy.run`` segment until interrupted or finished.

        The serve loop has already set ``_running`` (under the control lock),
        so this method must clear it on every exit path.
        """
        engine = self._ctx.policy.inference
        try:
            # Clear the local flag *before* checking the request flags: control
            # methods set their flag first and the segment-stop event second, so
            # a reset() or stop() racing with this start() is either seen here or
            # ends the freshly started loop on its first tick.
            self._segment_stop.clear()
            if (
                self._stop_requested.is_set()
                or self._reset_requested.is_set()
                or self._global_shutdown.is_set()
            ):
                return
            self._strategy.reset_control_state()
            self._emit(RolloutEvent.SEGMENT_STARTED)
            try:
                self._strategy.run(self._ctx)
            finally:
                engine.pause()
        finally:
            self._running.clear()
        if engine.failed:
            return  # the serve loop emits ENGINE_FAILED and shuts down
        if not (
            self._stop_requested.is_set() or self._reset_requested.is_set() or self._global_shutdown.is_set()
        ):
            self._emit(RolloutEvent.SEGMENT_ENDED)

    def _reset_robot(self) -> None:
        """Pause inference and return the robot home (the task was restored by :meth:`reset`)."""
        self._emit(RolloutEvent.RESET_STARTED)
        self._ctx.policy.inference.pause()
        if self._ctx.hardware.initial_position:
            self._strategy.return_to_initial_position(self._ctx.hardware)
            self._emit(RolloutEvent.RESET_DONE)
        else:
            logger.warning("No initial position captured — skipping the return move")
            self._emit(RolloutEvent.RESET_SKIPPED)

    def _emit(self, event: RolloutEvent) -> None:
        if self._on_event is None:
            return
        try:
            self._on_event(event)
        except Exception:  # a broken observer must not kill the serve loop
            logger.exception("Error in RolloutController event callback for %s", event)
