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

Interactive terminal controller for Unitree G1 locomotion + basic upper body (waist/arms).

Key features:
- Menu is generated from a single OPTIONS dictionary (easy to extend).
- Ctrl+C / Quit = StopMove only (keep current posture; NO prep/damp).
- Status prints LIVE joints (deg) from LowState.
- WALK/RUN teleop uses the "old working" keyboard control (pynput).
- PREP mode = stand + balance (NO Start/FSM200).
- READY/START mode = PREP + Start (FSM 200).
- StandUp/Squat/Sit/Lie-to-Stand use SetFsmId() fallback for compatibility.

Run:
  python3 g1_mode_controller.py --iface enp3s0
"""

from __future__ import annotations

import argparse
import sys
import time
from dataclasses import dataclass
from typing import Optional, Dict, Tuple, Callable, Any
from threading import Lock, Event
from collections import OrderedDict

# ---------- Unitree SDK2 ----------
from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelSubscriber
from unitree_sdk2py.g1.loco.g1_loco_client import LocoClient
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_  # type: ignore
from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient

# Optional teleop keyboard:
try:
    from pynput.keyboard import Listener, Key, KeyCode  # type: ignore
    HAVE_PYNPUT = True
except Exception:
    HAVE_PYNPUT = False


# ---------------- Joint map (for status) ----------------
JOINT_NAMES = {
    12: "WAIST_YAW",
    13: "WAIST_ROLL",
    14: "WAIST_PITCH",
    15: "L_SHOULDER_PITCH",
    16: "L_SHOULDER_ROLL",
    17: "L_SHOULDER_YAW",
    18: "L_ELBOW",
    19: "L_WRIST_ROLL",
    20: "L_WRIST_PITCH",
    21: "L_WRIST_YAW",
    22: "R_SHOULDER_PITCH",
    23: "R_SHOULDER_ROLL",
    24: "R_SHOULDER_YAW",
    25: "R_ELBOW",
    26: "R_WRIST_ROLL",
    27: "R_WRIST_PITCH",
    28: "R_WRIST_YAW",
}


def rad2deg(x: float) -> float:
    return float(x) * 180.0 / 3.141592653589793


# ---------------- Live State Monitor ----------------

class LiveStateMonitor:
    """Subscribes LowState and keeps last joint positions."""
    def __init__(self):
        self.lock = Lock()
        self.last_q_rad: Dict[int, float] = {}
        self.msg_count = 0
        self.last_rx_time = 0.0
        self.first_msg_evt = Event()

    def cb(self, msg: LowState_):
        motor_state = getattr(msg, "motor_state", None)
        if motor_state is None:
            return

        with self.lock:
            self.msg_count += 1
            self.last_rx_time = time.time()

            for idx in JOINT_NAMES.keys():
                try:
                    self.last_q_rad[idx] = float(motor_state[idx].q)
                except Exception:
                    pass

        self.first_msg_evt.set()

    def snapshot_deg(self) -> Tuple[int, float, Dict[int, float]]:
        with self.lock:
            age = (time.time() - self.last_rx_time) if self.last_rx_time else 1e9
            qdeg = {i: rad2deg(q) for i, q in self.last_q_rad.items()}
            return self.msg_count, age, qdeg


# ---------------- Upper body (waist/arms) publisher ----------------

@dataclass
class ArmRig:
    pub: object
    cmd: object
    crc: object
    waist_idx: int = 12

    def set_joint(self, idx: int, q: float, kp: float = 60.0, kd: float = 1.5, tau: float = 0.0):
        mc = self.cmd.motor_cmd[idx]
        mc.q = float(q)
        mc.dq = 0.0
        mc.tau = float(tau)
        mc.kp = float(kp)
        mc.kd = float(kd)

    def send(self):
        self.cmd.crc = self.crc.Crc(self.cmd)
        self.pub.Write(self.cmd)


def try_init_arm_sdk() -> Optional[ArmRig]:
    try:
        from unitree_sdk2py.core.channel import ChannelPublisher  # type: ignore
        from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_  # type: ignore
        from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_  # type: ignore
        from unitree_sdk2py.utils.crc import CRC  # type: ignore

        pub = ChannelPublisher("rt/arm_sdk", LowCmd_)
        pub.Init()

        cmd = unitree_hg_msg_dds__LowCmd_()
        crc = CRC()

        # enable arm_sdk mode
        cmd.motor_cmd[29].q = 1
        return ArmRig(pub=pub, cmd=cmd, crc=crc)
    except Exception as e:
        print("[arm_sdk] disabled:", e)
        return None


# ---------------- Core helpers ----------------

def init_loco(timeout: float = 10.0) -> LocoClient:
    bot = LocoClient()
    bot.SetTimeout(timeout)
    bot.Init()
    return bot


def init_msc(timeout: float = 5.0) -> MotionSwitcherClient:
    msc = MotionSwitcherClient()
    msc.SetTimeout(timeout)
    msc.Init()
    return msc


def safe_call(name: str, fn: Callable[..., Any], *args, **kwargs) -> Any:
    """
    Unitree python wrappers often return None on success.
    This helper prints [OK] on no-exception; prints [ERR] on exception.
    """
    try:
        ret = fn(*args, **kwargs)
        print(f"[OK] {name}: success (ret={ret})")
        return ret
    except Exception as e:
        code = getattr(e, "code", None)
        payload = getattr(e, "payload", None)
        print(f"[ERR] {name} failed: {repr(e)} code={code} payload={payload}")
        return None


def stop_only(bot: Optional[LocoClient]):
    """Stop walking/teleop velocity but do NOT change mode/posture."""
    if bot is None:
        return
    safe_call("StopMove", bot.StopMove)


def move_cmd(bot: LocoClient, vx: float, vy: float, om: float):
    """
    SDK variations exist. We try a safe call order.
    """
    try:
        bot.Move(vx, vy, om, True)
        return
    except Exception:
        pass

    try:
        bot.Move(vx, vy, om)
        return
    except Exception:
        pass

    try:
        bot.Move(vx, vy, om, continuous_move=True)  # type: ignore
        return
    except Exception as e:
        print("[ERR] Move failed in all signatures:", repr(e))


def fsm_set(bot: LocoClient, fsm_id: int, label: str):
    """
    Prefer direct methods if present; fallback to SetFsmId(fsm_id).
    """
    if hasattr(bot, label):
        fn = getattr(bot, label)
        return safe_call(label, fn)

    # common fallback in some sdk builds
    if hasattr(bot, "SetFsmId"):
        return safe_call(f"{label} (FSM {fsm_id}) via SetFsmId({fsm_id})", bot.SetFsmId, fsm_id)

    print(f"[ERR] Neither {label} nor SetFsmId exists in this SDK build.")
    return None


def balance_stand(bot: LocoClient, mode: int = 0):
    """
    Some builds require BalanceStand(balance_mode).
    Others expose SetBalanceMode(0/1).
    """
    if hasattr(bot, "BalanceStand"):
        try:
            return safe_call(f"BalanceStand({mode})", bot.BalanceStand, mode)
        except TypeError:
            return safe_call("BalanceStand()", bot.BalanceStand)

    if hasattr(bot, "SetBalanceMode"):
        return safe_call(f"SetBalanceMode({mode})", bot.SetBalanceMode, mode)

    print("[WARN] No BalanceStand/SetBalanceMode available in this SDK.")
    return None


# ---------------- PREP / READY sequences ----------------

def prep_mode(bot: LocoClient):
    """
    PREP = stand + balance (NO Start/FSM200).
    Matches your tested behavior:
      StopMove -> Damp -> StandUp(FSM4) -> ramp stand height -> BalanceStand(0) -> set final height.
    """
    print("[PREP] stand + balance (NO Start/FSM200) ...")
    safe_call("StopMove", bot.StopMove)
    safe_call("Damp", bot.Damp)

    # StandUp FSM 4
    fsm_set(bot, 4, "StandUp")

    # Stand height ramp (0.02 .. 0.50)
    if hasattr(bot, "SetStandHeight"):
        for h in [x / 100.0 for x in range(2, 51, 2)]:
            safe_call(f"SetStandHeight({h:.2f}m)", bot.SetStandHeight, float(h))
            time.sleep(0.03)

        # set a comfortable final height before balancing
        safe_call("SetStandHeight(0.22m)", bot.SetStandHeight, 0.22)
    else:
        print("[WARN] bot.SetStandHeight not available in this SDK build.")

    # Balance (static stand)
    balance_stand(bot, mode=0)

    # Re-send final height (common trick)
    if hasattr(bot, "SetStandHeight"):
        safe_call("SetStandHeight(0.22m)", bot.SetStandHeight, 0.22)

    print("[PREP] Done. (NO Start)")


def ready_start_mode(bot: LocoClient):
    """
    READY/START = PREP + Start (FSM200).
    """
    print("[READY] stand + balance + start (FSM 200) ...")
    prep_mode(bot)
    if hasattr(bot, "Start"):
        safe_call("Start (FSM 200)", bot.Start)
    else:
        # some builds: SetFsmId(200)
        fsm_set(bot, 200, "Start")
    print("[READY] Done. (FSM200 expected)")


# ---------------- Teleop (old keyboard function) ----------------

def teleop_loop(bot: LocoClient, speed_limit: float):
    if not HAVE_PYNPUT:
        print("[ERR] pynput not installed. Install: pip install pynput")
        return

    LIN_STEP = 0.05
    ANG_STEP = 0.2
    SEND_PERIOD = 0.10  # 10 Hz

    pressed = set()

    def on_press(k):
        if isinstance(k, KeyCode) and k.char:
            pressed.add(k.char.lower())
        else:
            pressed.add(k)

    def on_release(k):
        if isinstance(k, KeyCode) and k.char:
            pressed.discard(k.char.lower())
        else:
            pressed.discard(k)

    def key(name: str) -> bool:
        if name == "space":
            return Key.space in pressed
        if name == "esc":
            return Key.esc in pressed
        return name in pressed

    def clamp(v: float) -> float:
        return max(-speed_limit, min(speed_limit, v))

    vx = vy = om = 0.0
    last_send = 0.0

    print("\n--- TELEOP ---")
    print("Hold keys: W/S (vx), Q/E (vy), A/D (yaw), Space (stop).  ESC to exit.\n")

    listener = Listener(on_press=on_press, on_release=on_release)
    listener.start()

    try:
        while True:
            if key("esc"):
                stop_only(bot)
                break

            if key("w") and not key("s"):
                vx = clamp(vx + LIN_STEP)
            elif key("s") and not key("w"):
                vx = clamp(vx - LIN_STEP)
            else:
                vx = 0.0

            if key("q") and not key("e"):
                vy = clamp(vy + LIN_STEP)
            elif key("e") and not key("q"):
                vy = clamp(vy - LIN_STEP)
            else:
                vy = 0.0

            if key("a") and not key("d"):
                om = clamp(om + ANG_STEP)
            elif key("d") and not key("a"):
                om = clamp(om - ANG_STEP)
            else:
                om = 0.0

            if key("space"):
                vx = vy = om = 0.0

            now = time.time()
            if now - last_send >= SEND_PERIOD:
                # enable gait mode if supported (some SDKs require it for walking)
                if hasattr(bot, "SetBalanceMode"):
                    try:
                        bot.SetBalanceMode(1)
                    except Exception:
                        pass

                move_cmd(bot, vx, vy, om)
                last_send = now
                sys.stdout.write(f"\rvx={vx:+.2f}  vy={vy:+.2f}  yaw={om:+.2f}    ")
                sys.stdout.flush()

            time.sleep(0.005)

    except KeyboardInterrupt:
        stop_only(bot)
        print("\n--- teleop interrupted (StopMove only) ---\n")
    finally:
        listener.stop()
        print("\n--- teleop ended ---\n")


# ---------------- Presets for arms ----------------

LEFT_DOWN = [
    (12, 0.0),
    (15, +0.211),
    (16, +0.181),
    (17, -0.284),
    (18, +0.672),
    (19, -0.379),
    (20, -0.852),
    (21, -0.019),
]
RIGHT_DOWN = [
    (12, 0.0),
    (22, +0.087),
    (23, -0.271),
    (24, +0.323),
    (25, +0.691),
    (26, +0.240),
    (27, -0.771),
    (28, -0.176),
]

def apply_pose(arm: ArmRig, pose, kp: float = 60.0, kd: float = 1.5):
    for j, q in pose:
        arm.set_joint(j, q, kp=kp, kd=kd, tau=0.0)
    arm.send()


# ---------------- Main ----------------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--iface", required=True, help="NIC connected to G1 (example: enp3s0)")
    ap.add_argument("--timeout", type=float, default=10.0)
    ap.add_argument("--topic", default="", help="LowState topic override (default tries rt/lf/lowstate then rt/lowstate)")
    args = ap.parse_args()

    # DDS init ONCE
    ChannelFactoryInitialize(0, args.iface)

    # Live monitor
    monitor = LiveStateMonitor()
    topics = [args.topic] if args.topic else ["rt/lf/lowstate", "rt/lowstate"]
    sub = None
    last_err = None
    for t in topics:
        try:
            sub = ChannelSubscriber(t, LowState_)
            sub.Init(monitor.cb, 200)
            print(f"[OK] LowState subscribed: {t}")
            break
        except Exception as e:
            last_err = e
            print(f"[WARN] LowState subscribe failed for {t}: {e}", file=sys.stderr)
    if sub is None:
        print(f"[WARN] Could not subscribe LowState. Last error: {last_err}")

    # Init clients
    print("Initializing LocoClient...")
    bot: Optional[LocoClient] = None
    try:
        bot = init_loco(timeout=args.timeout)
    except Exception as e:
        print("[ERR] Loco init failed:", e)

    print("Initializing MotionSwitcherClient...")
    msc: Optional[MotionSwitcherClient] = None
    try:
        msc = init_msc(timeout=5.0)
        status, result = msc.CheckMode()
        print(f"[MSC] Current mode: {result.get('name') or '(none)'}")
    except Exception as e:
        print("[WARN] MotionSwitcher init failed:", e)

    arm = try_init_arm_sdk()

    # Snapshot on start
    if sub is not None and monitor.first_msg_evt.wait(timeout=2.0):
        cnt, age, qdeg = monitor.snapshot_deg()
        print("\n[STATE] Current joint snapshot (deg):")
        print(f"  lowstate_msgs={cnt}  age={age:.2f}s")
        for idx in sorted(qdeg.keys()):
            print(f"  [{idx:02d}] {JOINT_NAMES.get(idx,'?'):<18} {qdeg[idx]:+8.2f}°")
        print()

    def need_bot() -> bool:
        if bot is None:
            print("[ERR] No LocoClient available.")
            return False
        return True

    def need_arm() -> bool:
        if arm is None:
            print("[ERR] arm_sdk not available.")
            return False
        return True

    # ----- Dictionary-driven menu -----
    OPTIONS: "OrderedDict[int, Dict[str, Any]]" = OrderedDict()

    OPTIONS[1] = {"label": "PREP mode (stand + balance)  [NO start/FSM200]",
                  "fn": lambda: prep_mode(bot) if need_bot() else None}

    OPTIONS[2] = {"label": "READY/START mode (stand + balance + start, FSM 200)",
                  "fn": lambda: ready_start_mode(bot) if need_bot() else None}

    OPTIONS[3] = {"label": "ZeroTorque (FSM 0)",
                  "fn": lambda: fsm_set(bot, 0, "ZeroTorque") if need_bot() else None}

    OPTIONS[4] = {"label": "Damp (FSM 1)",
                  "fn": lambda: fsm_set(bot, 1, "Damp") if need_bot() else None}

    OPTIONS[5] = {"label": "StandUp (FSM 4)",
                  "fn": lambda: fsm_set(bot, 4, "StandUp") if need_bot() else None}

    OPTIONS[6] = {"label": "Squat (FSM 2)",
                  "fn": lambda: fsm_set(bot, 2, "Squat") if need_bot() else None}

    OPTIONS[7] = {"label": "Sit (FSM 3)",
                  "fn": lambda: fsm_set(bot, 3, "Sit") if need_bot() else None}

    OPTIONS[8] = {"label": "LowStand",
                  "fn": lambda: safe_call("LowStand", bot.LowStand) if (need_bot() and hasattr(bot, "LowStand")) else print("[ERR] LowStand not available")}

    OPTIONS[9] = {"label": "HighStand",
                  "fn": lambda: safe_call("HighStand", bot.HighStand) if (need_bot() and hasattr(bot, "HighStand")) else print("[ERR] HighStand not available")}

    OPTIONS[10] = {"label": "Lie-to-Stand (FSM 702) (if supported)",
                   "fn": lambda: fsm_set(bot, 702, "Lie2StandUp") if need_bot() else None}

    OPTIONS[11] = {"label": "Walk teleop (keyboard: W/S Q/E A/D, ESC exit)",
                   "fn": lambda: teleop_loop(bot, speed_limit=0.6) if need_bot() else None}

    OPTIONS[12] = {"label": "Run teleop (keyboard: W/S Q/E A/D, ESC exit)",
                   "fn": lambda: teleop_loop(bot, speed_limit=1.2) if need_bot() else None}

    OPTIONS[16] = {"label": "Status (LIVE joints)",
                   "fn": lambda: show_status(bot, arm, msc, sub, monitor)}

    OPTIONS[17] = {"label": "MotionSwitcher: Select mode 'ai' (supported)",
                   "fn": lambda: msc_select_ai(msc)}

    OPTIONS[18] = {"label": "MotionSwitcher: Release current mode",
                   "fn": lambda: msc_release(msc)}

    OPTIONS[19] = {"label": "MotionSwitcher: Show current mode",
                   "fn": lambda: msc_show(msc)}

    OPTIONS[20] = {"label": "Seated mode ON (Sit)",
                   "fn": lambda: fsm_set(bot, 3, "Sit") if need_bot() else None}

    OPTIONS[21] = {"label": "Seated mode OFF (StandUp + PREP)",
                   "fn": lambda: seated_off(bot) if need_bot() else None}

    OPTIONS[22] = {"label": "Reconnect (re-init Loco + MotionSwitcher)",
                   "fn": lambda: reconnect(args, lambda b: set_bot_ref(b), lambda m: set_msc_ref(m))}

    # Helpers need access to outer bot/msc refs:
    def set_bot_ref(new_bot: Optional[LocoClient]):
        nonlocal bot
        bot = new_bot

    def set_msc_ref(new_msc: Optional[MotionSwitcherClient]):
        nonlocal msc
        msc = new_msc


    def print_menu():
        print("\n" + "=" * 70)
        print("G1 MODE CONTROLLER - select option:")
        print("=" * 70)
        for k, v in OPTIONS.items():
            print(f"{k:>3}) {v['label']}")
        print("  0) Quit (StopMove only)")
        print("=" * 70)

    print_menu()

    while True:
        try:
            choice = input("\nSelect option (0-22): ").strip()
        except KeyboardInterrupt:
            stop_only(bot)
            print("\n[EXIT] Ctrl+C → StopMove only. Keeping current posture. Bye.")
            break
        except EOFError:
            stop_only(bot)
            print("\n[EXIT] EOF → StopMove only. Keeping current posture. Bye.")
            break

        if not choice:
            continue

        if choice == "0":
            stop_only(bot)
            print("\n[EXIT] Quit → StopMove only. Keeping current posture. Bye.")
            break

        try:
            opt = int(choice)
        except ValueError:
            print("Please enter a number.")
            continue

        if opt not in OPTIONS:
            print("Unknown option.")
            print_menu()
            continue

        try:
            OPTIONS[opt]["fn"]()
        except KeyboardInterrupt:
            stop_only(bot)
            print("\n[INTERRUPT] Ctrl+C → StopMove only. Keeping current posture.")
        except Exception as e:
            print("[ERR] Action failed:", repr(e))

        print_menu()


# --------- extra action helpers that need current refs ---------

def show_status(bot, arm, msc, sub, monitor):
    print("Status:")
    print(f"  LocoClient:      {'OK' if bot is not None else 'NOT INITIALISED'}")
    print(f"  MotionSwitcher:  {'OK' if msc is not None else 'DISABLED'}")
    if msc is not None:
        try:
            _s, r = msc.CheckMode()
            print(f"  MSC mode:        {r.get('name') or '(none)'}")
        except Exception:
            print("  MSC mode:        (error reading)")
    print(f"  arm_sdk:         {'OK' if arm is not None else 'DISABLED'}")
    print(f"  pynput:          {'OK' if HAVE_PYNPUT else 'NOT INSTALLED'}")

    if sub is None:
        print("  LowState:        NOT SUBSCRIBED")
        return

    cnt, age, qdeg = monitor.snapshot_deg()
    print(f"  LowState:        OK  msgs={cnt}  age={age:.2f}s")
    print("  Current joints (deg):")
    for idx in sorted(qdeg.keys()):
        print(f"    [{idx:02d}] {JOINT_NAMES.get(idx,'?'):<18} {qdeg[idx]:+8.2f}°")


def msc_select_ai(msc: Optional[MotionSwitcherClient]):
    if msc is None:
        print("[ERR] MotionSwitcher not available.")
        return
    ret = safe_call("SelectMode('ai')", msc.SelectMode, "ai")
    # ret could be (7004,None) etc; we just print it from safe_call.


def msc_release(msc: Optional[MotionSwitcherClient]):
    if msc is None:
        print("[ERR] MotionSwitcher not available.")
        return
    safe_call("ReleaseMode()", msc.ReleaseMode)


def msc_show(msc: Optional[MotionSwitcherClient]):
    if msc is None:
        print("[ERR] MotionSwitcher not available.")
        return
    try:
        s, r = msc.CheckMode()
        print(f"[MSC] Current mode: {r.get('name') or '(none)'}")
    except Exception as e:
        print("[ERR] CheckMode failed:", repr(e))


def seated_off(bot: LocoClient):
    print("[SEATED] Leaving seated: StandUp + PREP (no Start)...")
    fsm_set(bot, 4, "StandUp")
    prep_mode(bot)


def reconnect(args, set_bot, set_msc):
    print("Reconnecting LocoClient...")
    try:
        b = init_loco(timeout=args.timeout)
        set_bot(b)
        print("[OK] LocoClient reconnected.")
    except Exception as e:
        print("[ERR] Loco reconnect failed:", repr(e))
        set_bot(None)

    print("Reconnecting MotionSwitcherClient...")
    try:
        m = init_msc(timeout=5.0)
        set_msc(m)
        try:
            _s, r = m.CheckMode()
            print(f"[MSC] Current mode: {r.get('name') or '(none)'}")
        except Exception:
            pass
        print("[OK] MotionSwitcher reconnected.")
    except Exception as e:
        print("[ERR] MotionSwitcher reconnect failed:", repr(e))
        set_msc(None)


if __name__ == "__main__":
    main()