#!/usr/bin/env python3
"""Switch the Flexiv active tool (TCP + payload) to a Studio-configured name.

Must run under RDK Python (rdk_py310). Robot must be IDLE; stop run_real.sh first
if ros2_control already owns the RDK connection.

  /home/cc/rdk_py310/bin/python ~/ros2/scripts/rdk_switch_tool.py --list
  /home/cc/rdk_py310/bin/python ~/ros2/scripts/rdk_switch_tool.py --name Flange
  /home/cc/rdk_py310/bin/python ~/ros2/scripts/rdk_switch_tool.py   # site.yaml robot.tool_name
"""

from __future__ import annotations

import argparse
import sys
import time


def _single_arm_groups(robot):
    """RDK 1.9 Tool API is single-arm (no JointGroup). Older builds use groups."""
    info = robot.info()
    groups = getattr(info, "single_arm_groups", None) or {}
    if isinstance(groups, dict):
        keys = list(groups.keys())
    else:
        keys = list(groups)
    if keys:
        return keys
    try:
        import flexivrdk

        if hasattr(flexivrdk, "JointGroup"):
            return [flexivrdk.JointGroup.ARM_1]
    except Exception:
        pass
    return [None]  # sentinel: call Tool.name()/Switch() without a group


def _tool_name(tool, group):
    if group is None:
        return tool.name()
    return tool.name(group)


def _tool_params(tool, group):
    if group is None:
        return tool.params()
    return tool.params(group)


def _tool_switch(tool, group, target: str) -> None:
    if group is None:
        tool.Switch(target)
        return
    tool.Switch(group, target)


def _group_name(group) -> str:
    if group is None:
        return "arm"
    try:
        import flexivrdk

        return str(flexivrdk.kJointGroupNames[group])
    except Exception:
        return str(group)


def _fmt_params(params) -> str:
    mass = getattr(params, "mass", None)
    com = list(getattr(params, "CoM", []) or [])
    tcp = list(getattr(params, "tcp_location", []) or [])
    com_s = ",".join(f"{float(v):.4f}" for v in com[:3]) if com else "?"
    tcp_s = ",".join(f"{float(v):.4f}" for v in tcp[:3]) if tcp else "?"
    return f"mass={mass} kg  CoM=[{com_s}] m  TCP_xyz=[{tcp_s}] m"


def _ensure_idle(robot, mode) -> None:
    try:
        robot.SwitchMode(mode.IDLE)
        return
    except Exception:
        pass
    if hasattr(robot, "Enable"):
        robot.Enable()
    else:
        robot.ServoOn()
    for _ in range(40):
        if robot.operational():
            break
        time.sleep(0.25)
    if not robot.operational():
        raise RuntimeError(
            "robot not operational — AUTO Remote + Enable, release E-stop"
        )
    robot.SwitchMode(mode.IDLE)


def _resolve_name(tool, want: str) -> str:
    want = (want or "").strip()
    if not want:
        return ""
    names = list(tool.list() or [])
    if want in names:
        return want
    lower = {n.lower(): n for n in names}
    hit = lower.get(want.lower())
    if hit:
        return hit
    raise SystemExit(
        "tool %r not found in Studio pool.\n  available: %s"
        % (want, ", ".join(names) or "(empty)")
    )


def _iface_ipv4s(name: str) -> list:
    """FastDDS whitelist wants IPv4 addresses, not NIC names like enp4s0."""
    import socket
    import fcntl
    import struct

    name = (name or "").strip()
    if not name:
        return []
    if name.count(".") == 3:
        return [name]
    addrs = []
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        ifreq = struct.pack("256s", name.encode("utf-8")[:15])
        addrs.append(socket.inet_ntoa(fcntl.ioctl(s.fileno(), 0x8915, ifreq)[20:24]))
        s.close()
    except OSError:
        pass
    try:
        import subprocess

        out = subprocess.check_output(
            ["ip", "-4", "-o", "addr", "show", "dev", name],
            text=True,
            stderr=subprocess.DEVNULL,
        )
        for line in out.splitlines():
            parts = line.split()
            if "inet" in parts:
                cidr = parts[parts.index("inet") + 1]
                addrs.append(cidr.split("/")[0])
    except Exception:
        pass
    # unique, keep order
    seen = set()
    uniq = []
    for a in addrs:
        if a and a not in seen:
            seen.add(a)
            uniq.append(a)
    return uniq


def _rdk_whitelist(iface: str) -> list[list[str]]:
    """RDK 手册要的是「连机器人的那张网卡」。enp4s0 上还有办公网 IP，不能整张卡的地址都塞进白名单。"""
    import os

    lists: list[list[str]] = []
    pc = (os.environ.get("SITE_PC_IP") or "").strip()
    if pc:
        lists.append([pc])
    if iface:
        lists.append([iface])
        robot_lan = [a for a in _iface_ipv4s(iface) if a.startswith("192.168.2.")]
        if robot_lan:
            lists.append(robot_lan)
    # unique
    seen: set[tuple[str, ...]] = set()
    out: list[list[str]] = []
    for item in lists:
        key = tuple(item)
        if key not in seen:
            seen.add(key)
            out.append(item)
    return out


def connect_rdk(sn: str, iface: str = ""):
    import flexivrdk

    for listen in _rdk_whitelist(iface):
        print("RDK whitelist:", ", ".join(listen))
        for args in ((sn, listen), (sn, listen, True), (sn, listen, True, False)):
            try:
                return flexivrdk.Robot(*args)
            except TypeError:
                continue
            except RuntimeError as exc:
                msg = str(exc)
                if "whitelist" in msg.lower() or "filtered out" in msg.lower():
                    print("  whitelist rejected, try next")
                    break
                if "Failed to connect" in msg or "failed to connect" in msg.lower():
                    print("  search timeout with this whitelist")
                    break
                raise
    print("retry without iface filter")
    return flexivrdk.Robot(sn)


def main() -> int:
    ap = argparse.ArgumentParser(description="Switch Flexiv Studio tool / TCP")
    ap.add_argument("--sn", default="", help="robot serial (default: site.yaml)")
    ap.add_argument("--iface", default="", help="NIC whitelist for RDK (e.g. enp4s0)")
    ap.add_argument(
        "--name",
        default="",
        help="Studio tool name (Flange = no tool). Empty = site.yaml robot.tool_name",
    )
    ap.add_argument("--list", action="store_true", help="list tools and exit")
    args = ap.parse_args()

    sn = (args.sn or "").strip()
    want = (args.name or "").strip()
    if not sn:
        print("ERROR: --sn is required (run via run_real.sh / rdk_switch_tool.sh)", file=sys.stderr)
        return 2

    import flexivrdk

    mode = flexivrdk.Mode
    iface = (args.iface or "").strip()
    print(f"=== RDK tool  SN={sn}  iface={iface or 'auto'} ===")
    robot = None
    try:
        robot = connect_rdk(sn, iface)
        if robot.fault():
            print("fault present, ClearFault ...")
            if not robot.ClearFault():
                print("ERROR: fault cannot be cleared", file=sys.stderr)
                return 1
        tool = flexivrdk.Tool(robot)
        names = list(tool.list() or [])
        groups = _single_arm_groups(robot)
        print("Studio tools:", ", ".join(names) or "(none)")
        for g in groups:
            cur = _tool_name(tool, g)
            extra = ""
            try:
                extra = "  " + _fmt_params(_tool_params(tool, g))
            except Exception:
                pass
            print(f"  active [{_group_name(g)}]: {cur}{extra}")

        if args.list:
            return 0
        if not want:
            print("tool_name empty — keep current (set robot.tool_name in site.yaml to switch)")
            print("RDK OK")
            return 0

        target = _resolve_name(tool, want)
        already = all(_tool_name(tool, g) == target for g in groups)
        if already:
            print(f"already on [{target}]")
            print("RDK OK")
            return 0

        print(f"switching → [{target}] (requires IDLE)")
        _ensure_idle(robot, mode)
        for g in groups:
            _tool_switch(tool, g, target)
        for g in groups:
            cur = _tool_name(tool, g)
            extra = ""
            try:
                extra = "  " + _fmt_params(_tool_params(tool, g))
            except Exception:
                pass
            print(f"  now [{_group_name(g)}]: {cur}{extra}")
            if cur != target:
                print(f"ERROR: switch failed, still {cur!r}", file=sys.stderr)
                return 1
        print("RDK OK")
        return 0
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        print("  Stop run_real.sh first (RDK is exclusive).", file=sys.stderr)
        print("  Pendant: AUTO Remote + Enable Ethernet/RDK.", file=sys.stderr)
        return 1
    finally:
        robot = None


if __name__ == "__main__":
    raise SystemExit(main())