#!/usr/bin/env python3
"""
R1 arm-action RPC client — the R1-NATIVE way to move the upper body (arms + head)
while the firmware keeps balancing the legs.

Why this and not rt/arm_sdk: the R1 firmware does NOT expose live joint streaming
of the head/arms. The 'arm' service (a.k.a. r1_arm_example) owns rt/arm_sdk
internally; clients drive the upper body ONLY through this RPC — preset gesture
actions (by id) and custom "teach" actions (by name, recorded via the app's
teaching). Captured from the app's own traffic (2026-07-06):
  release_arm -> ExecuteAction(99); replay recorded motion -> ExecuteCustomAction(name).

Service "arm", api version "1.0.0.14". Same Client base as R1LocoClient.
State is published on rt/arm/action/state: {"holding":bool,"id":int,"name":str}.
"""
import json
from unitree_sdk2py.rpc.client import Client

ARM_ACTION_SERVICE_NAME = "arm"
ARM_ACTION_API_VERSION = "1.0.0.14"

API_EXECUTE_ACTION = 7106          # {"action_id": N}   preset actions
API_GET_ACTION_LIST = 7107         # -> JSON list
API_EXECUTE_CUSTOM_ACTION = 7108   # {"action_name": "..."}  recorded teach actions
API_RECORD = 7110                  # start = {"action_name":NAME} ; stop+save = "" (empty str)
API_STOP_CUSTOM_ACTION = 7113

RELEASE_ARM = 99                   # go compliant / release a held action


class R1ArmActionClient(Client):
    def __init__(self):
        super().__init__(ARM_ACTION_SERVICE_NAME, False)

    def Init(self):
        self._SetApiVerson(ARM_ACTION_API_VERSION)  # SDK spells it "Verson"
        for api in (API_EXECUTE_ACTION, API_GET_ACTION_LIST,
                    API_EXECUTE_CUSTOM_ACTION, API_RECORD, API_STOP_CUSTOM_ACTION):
            self._RegistApi(api, 0)

    def ExecuteAction(self, action_id: int):
        """Run a preset action by id (99 = release). Returns rc (0 = ok)."""
        code, _ = self._Call(API_EXECUTE_ACTION, json.dumps({"action_id": int(action_id)}))
        return code

    def ExecuteCustomAction(self, action_name: str):
        """Replay a recorded teach action by name (from the app's teaching)."""
        code, _ = self._Call(API_EXECUTE_CUSTOM_ACTION, json.dumps({"action_name": str(action_name)}))
        return code

    def StopCustomAction(self):
        code, _ = self._Call(API_STOP_CUSTOM_ACTION, "{}")
        return code

    def ReleaseArm(self):
        return self.ExecuteAction(RELEASE_ARM)

    def GetActionList(self):
        """Return (rc, data) — data is the JSON list of available actions."""
        code, data = self._Call(API_GET_ACTION_LIST, "")
        parsed = None
        if code == 0 and data:
            try:
                parsed = json.loads(data)
            except (ValueError, TypeError):
                parsed = data
        return code, parsed