#!/usr/bin/env python3
"""
g1_arm_actions_cli.py

Run Unitree G1 Arm ActionClient actions directly on the robot from CLI.

Changes from your request:
- Network interface is OPTIONAL (no longer required).
  If you pass it, it will use it; otherwise it uses default DDS interface.

Usage:
  # list actions
  python3 g1_arm_actions_cli.py list

  # run by id
  python3 g1_arm_actions_cli.py 6

  # run by name (quotes recommended)
  python3 g1_arm_actions_cli.py "face wave"

  # specify interface explicitly (optional)
  python3 g1_arm_actions_cli.py "shake hand" --iface enp3s0

  # force auto release
  python3 g1_arm_actions_cli.py "hug" --auto-release --sleep 2
"""

import argparse
import sys
import time
from dataclasses import dataclass
from typing import Optional

from unitree_sdk2py.core.channel import ChannelFactoryInitialize
from unitree_sdk2py.g1.arm.g1_arm_action_client import G1ArmActionClient, action_map


@dataclass
class TestOption:
    name: str
    id: int


OPTION_LIST = [
    TestOption(name="release arm", id=0),
    TestOption(name="shake hand", id=1),
    TestOption(name="high five", id=2),
    TestOption(name="hug", id=3),
    TestOption(name="high wave", id=4),
    TestOption(name="clap", id=5),
    TestOption(name="face wave", id=6),
    TestOption(name="left kiss", id=7),
    TestOption(name="heart", id=8),
    TestOption(name="right heart", id=9),
    TestOption(name="hands up", id=10),
    TestOption(name="x-ray", id=11),
    TestOption(name="right hand up", id=12),
    TestOption(name="reject", id=13),
    TestOption(name="right kiss", id=14),
    TestOption(name="two-hand kiss", id=15),
]

# Actions that your original script auto-released after ~2s
DEFAULT_AUTO_RELEASE_IDS = {1, 2, 3, 8, 9, 10, 11, 12, 13}


def find_option(query: str) -> Optional[TestOption]:
    q = query.strip().lower()

    # try id
    try:
        as_int = int(q)
        for opt in OPTION_LIST:
            if opt.id == as_int:
                return opt
    except ValueError:
        pass

    # try name
    for opt in OPTION_LIST:
        if opt.name.lower() == q:
            return opt

    return None


def print_list():
    print("\nAvailable actions:")
    for opt in OPTION_LIST:
        print(f"  {opt.id:>2}  {opt.name}")
    print("")


def execute_action(client: G1ArmActionClient, opt: TestOption):
    if opt.name not in action_map:
        raise RuntimeError(f"Action '{opt.name}' not found in action_map. Check SDK version.")
    print(f"[RUN] {opt.name} (id={opt.id})")
    client.ExecuteAction(action_map.get(opt.name))


def init_dds(iface: Optional[str]):
    # If iface provided, use it; otherwise default interface.
    if iface:
        ChannelFactoryInitialize(0, iface)
    else:
        ChannelFactoryInitialize(0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument(
        "action",
        help='Action id/name, or "list". Examples: 6  |  "face wave"  |  list',
    )
    ap.add_argument("--iface", default=None, help="Optional network interface, e.g. enp3s0/eth0/wlan0")
    ap.add_argument("--timeout", type=float, default=10.0, help="Action client timeout (seconds)")
    ap.add_argument(
        "--auto-release",
        action="store_true",
        help="After action, run 'release arm' automatically (like your demo for some actions).",
    )
    ap.add_argument(
        "--sleep",
        type=float,
        default=2.0,
        help="Seconds to wait before auto-release (only used with auto-release).",
    )
    ap.add_argument(
        "--no-prompt",
        action="store_true",
        help="Skip the 'Press Enter to continue' safety prompt.",
    )
    args = ap.parse_args()

    if args.action.strip().lower() == "list":
        print_list()
        return

    opt = find_option(args.action)
    if not opt:
        print(f"[ERR] No matching action for: {args.action!r}")
        print_list()
        sys.exit(2)

    print("WARNING: Please ensure there are no obstacles around the robot while running this script.")
    if not args.no_prompt:
        input("Press Enter to continue... ")

    # DDS init (optional iface)
    init_dds(args.iface)

    # Arm action client
    client = G1ArmActionClient()
    client.SetTimeout(args.timeout)
    client.Init()

    execute_action(client, opt)

    # Decide auto-release behavior
    do_auto_release = args.auto_release or (opt.id in DEFAULT_AUTO_RELEASE_IDS)

    if do_auto_release and opt.id != 0:
        time.sleep(max(0.0, args.sleep))
        print("[AUTO] release arm")
        client.ExecuteAction(action_map.get("release arm"))

    time.sleep(0.5)
    print("[DONE]")


if __name__ == "__main__":
    main()