# SPDX-FileCopyrightText: Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# SPDX-License-Identifier: BSD-3-Clause

"""Single-arm XR teleoperation agent."""

from __future__ import annotations

import torch

from .base_agent import BaseTeleopAgent
from .xr_controller import XrTeleopController
from .frame_utils import convert_action_world_to_root


class SingleArmXrAgent(BaseTeleopAgent):
    """Single-arm teleoperation: one XR controller -> one robot arm + gripper.

    Action: arm(7) | gripper(1) = 8D
    """

    def __init__(self, env, simulation_app, *,
                 arm: str = "left",
                 pos_sensitivity: float = 1.0,
                 rot_sensitivity: float = 0.3,
                 xr_client=None,
                 debug_viewports: bool = True):
        super().__init__(env, simulation_app, debug_viewports=debug_viewports)

        self.xr_client = xr_client
        self.teleop = XrTeleopController(
            arm=arm, pos_sensitivity=pos_sensitivity,
            rot_sensitivity=rot_sensitivity, xr_client=xr_client,
        )
        self.teleop.add_callback("RESET", self.request_reset)
        self._last_root = None

    def _ik_action_term_names(self) -> list[str]:
        return ["arm_action"]

    def on_reset(self):
        self.teleop.reset()
        self._last_root = None

    def assemble_action(self, obs) -> torch.Tensor:
        policy_obs = obs["policy"]
        eef_pos = policy_obs["eef_pos"][0].cpu().numpy()
        eef_quat = policy_obs["eef_quat"][0].cpu().numpy()

        action = self.teleop.advance(current_eef_pos=eef_pos, current_eef_quat=eef_quat)

        robot = self.env.unwrapped.scene["robot"]
        if self.teleop.grip_active or self._last_root is None:
            self._last_root = convert_action_world_to_root(action, robot)[:7].clone()

        # arm(7) | gripper(1)
        return torch.cat([self._last_root, action[7:8]])

    def cleanup(self):
        self.teleop.close()
        super().cleanup()