"""Base v6 agent for any env that exposes a ``robot`` capability.

Subclass :class:`RobotAgent`, set ``self.model`` and ``self.adapter`` in
``__init__``, and the base owns the rest.

The base calls the adapter and model at the right moments::

    setup_robot      -> adapter.bind(spaces)       # once after connect
    on_episode_start -> adapter.reset()            # per episode; model is stateless
    select_action    -> adapt_observation -> model.ainfer -> pop chunk -> adapt_action

``model.ainfer`` always returns a ``[T, A]`` chunk; :meth:`RobotAgent.select_action`
executes it open-loop, re-inferring only once the active chunk is spent.

Most policies use :class:`~hud.agents.robot.adapter.LeRobotAdapter`; a policy whose
spaces match the env natively can set ``adapter = None`` (raw pass-through).
"""

from __future__ import annotations

from collections import deque
from typing import TYPE_CHECKING, Any, ClassVar

import numpy as np

from hud.agents.base import Agent
from hud.capabilities.robot import RobotClient

from .record import Recorder

if TYPE_CHECKING:
    from hud.eval.run import Run

    from ._types import ActionArray
    from .adapter import Adapter
    from .model import Model

ROBOT_PROTOCOL = "openpi/0"


class RobotAgent(Agent):
    """Drive a ``robot`` side-channel for one :class:`~hud.client.Run`.

    **Subclass contract:** in ``__init__`` set ``self.model`` (a
    :class:`~hud.agents.robot.model.Model`) and ``self.adapter`` (an
    :class:`~hud.agents.robot.adapter.Adapter`, or ``None`` for raw pass-through).

    **Override if needed:**

    - :attr:`robot_protocol` — class attr if not ``openpi/0``
    - :meth:`on_episode_start` — mostly internal; override (with ``super()``) to
      add per-episode setup (e.g. reading the env contract).
    - :meth:`should_stop` — custom early-exit condition beyond ``obs["terminated"]``
    - :meth:`select_action` — only for a wholly different inference path
    - :attr:`log_every` — class-level print frequency (0 = off)
    """

    robot_protocol: ClassVar[str] = ROBOT_PROTOCOL
    #: How often (in steps) to print a step-progress line. 0 = off.
    log_every: ClassVar[int] = 20
    #: Opt-in: also save a LeRobot v3 dataset of every (obs, action) pair to disk
    #: (the ``--save`` flag). Telemetry streams regardless; see :mod:`.record`.
    save: bool = False

    #: Runs the policy (preprocess → forward → postprocess). Subclasses set this.
    model: Model | None = None
    #: Translates env<->policy spaces. Subclasses set this; ``None`` = raw pass-through.
    adapter: Adapter | None = None

    _prompt: str = ""
    #: The env's action / observation contract features (from ``client.spaces()``),
    #: named ``_env_*`` to mark them as env-side values (not the policy's spaces).
    _env_action_space: dict[str, Any]
    _env_obs_space: dict[str, Any]
    #: Unexecuted tail of the current policy chunk; popped one action per step.
    _active_chunk: deque[ActionArray]
    #: Control-tick index, incremented per executed action.
    _tick: int
    #: Records all telemetry (observation/inference steps + video) and, when ``save``, a
    #: LeRobot dataset. Agent-lifetime (the dataset spans every episode); created lazily.
    _recorder: Recorder | None = None

    def setup_robot(self, client: RobotClient) -> None:
        """Discover the env's action/observation layout and bind the adapter to it."""
        self._env_action_space, self._env_obs_space = client.spaces()
        if self.adapter is not None:
            self.adapter.bind(self._env_action_space, self._env_obs_space)

    def on_episode_start(self, run: Run, client: RobotClient, *, prompt: str) -> None:
        """Store the prompt and reset per-episode state before the act loop.

        The model is stateless (per-episode state lives here, not on the shared model), so
        only the adapter is reset. Override (calling ``super()`` first) for extra setup.
        """
        self._prompt = prompt
        self._active_chunk = deque()
        self._tick = 0
        # One recorder for the agent's life so its LeRobot dataset spans every episode;
        # begin() opens this episode (fresh video stream, prompt) and takes the run it records onto.
        if self._recorder is None:
            self._recorder = Recorder(client, save=self.save)
        self._recorder.begin(run, prompt)
        if self.adapter is not None:
            self.adapter.reset()

    def should_stop(self, obs: dict[str, Any], *, step: int, max_steps: int) -> bool:
        """Return True to break out of the step loop (before ``select_action``)."""
        return bool(obs.get("terminated"))

    async def select_action(self, obs: dict[str, Any]) -> ActionArray:
        """Pop the next action, re-inferring a ``[T, A]`` chunk once the active one is
        spent, then adapt it to env space. Override only for a different inference path.
        """
        if self.model is None:
            raise RuntimeError(f"{type(self).__name__} must set self.model in __init__")
        if not self._active_chunk:
            batch = (
                obs if self.adapter is None else self.adapter.adapt_observation(obs, self._prompt)
            )
            chunk = np.atleast_2d(await self.model.ainfer(batch))  # [T, A]
            self._active_chunk = deque(chunk)
            assert self._recorder is not None  # set in on_episode_start
            self._recorder.record_inference(chunk, tick=self._tick)
        self._tick += 1
        raw = self._active_chunk.popleft()
        return raw if self.adapter is None else self.adapter.adapt_action(raw, obs)

    async def __call__(self, run: Run, *, max_steps: int | None = None) -> None:
        step_limit = max_steps if max_steps is not None else int(getattr(self, "max_steps", 520))
        cap = run.client.binding(self.robot_protocol)
        client = await RobotClient.connect(cap)
        try:
            self.setup_robot(client)
            prompt = run.prompt
            if not isinstance(prompt, str):
                raise TypeError(
                    f"run.prompt must be a str, got {type(prompt).__name__}: {prompt!r}"
                )
            self.on_episode_start(run, client, prompt=prompt)
            print(f"[agent] episode started: {prompt!r} (max_steps={step_limit})", flush=True)

            assert self._recorder is not None  # set in on_episode_start above
            for step in range(step_limit):
                obs = await client.get_observation()
                self._recorder.record_observation(obs, tick=step)

                if self.should_stop(obs, step=step, max_steps=step_limit):
                    print(f"[agent] env reported terminated at step {step}", flush=True)
                    break

                action = await self.select_action(obs)
                self._recorder.record_action(action)
                await client.send_action(action)

                if self.log_every and step % self.log_every == 0:
                    preview = np.array2string(action, precision=3, suppress_small=True)
                    print(f"[agent] step {step}/{step_limit} action={preview}", flush=True)
            else:
                print(f"[agent] reached max_steps={step_limit}", flush=True)

            run.trace.status = "completed"
            run.trace.content = "done"
        finally:
            if self._recorder is not None:
                self._recorder.end()  # flush video tails + commit the LeRobot episode
            await client.close()


__all__ = ["ROBOT_PROTOCOL", "RobotAgent"]
