#!/usr/bin/env python3
"""RDK 控制桥（必须在 rdk_py310 / Python 3.10 下运行）。

为什么需要这个进程
------------------
非夕的 `flexivrdk` Python 绑定与 ROS 环境的 Python 3.11 ABI 不兼容，且 RDK 对
同一台机器人是**独占访问**：多个进程同时连同一个 SN 会失败。所以框架里所有
RDK 操作都收敛到本进程，由它串行化，业务节点通过 JSON 行协议远程调用。

这也让 `arm_driver_flexiv` 不必 import flexivrdk，符合
`docs/INTERFACES.md` 的约束（业务节点不得 import flexivrdk）。

软拖拽为什么用 NRT_JOINT_IMPEDANCE
----------------------------------
AUTO(Remote) 模式下示教器上的 Free Drive 按钮不可用（控制权在 Remote）。
RDK 的 FloatingJoint/FloatingCartesian 通常需要额外授权。
所以默认用低刚度 NRT_JOINT_IMPEDANCE：目标跟踪实测关节角 + 低刚度，
于是手臂可以用手推动，同时 Remote 保持开启。

线程安全
--------
`flexivrdk.Robot` 不是线程安全的。所有 robot API 调用（含拖拽循环）共用同一
把锁 `rlock`。拖拽期间的关节读取走缓存，避免 Enter 与 Stop() 竞争。

行协议（JSON 只写 stdout；RDK 自身日志走 stderr）
-------------------------------------------------
    {"cmd":"ping"}
    {"cmd":"connect","sn":"<site.yaml 的 robot.sn>","iface":"<robot.iface>"}
    {"cmd":"enable"}                              # 使能 + 清故障（幂等）
    {"cmd":"clear_fault"}                         # 显式清故障并使能
    {"cmd":"switch_tool","name":"tool1"}          # 切工具；name 空=只查询
    {"cmd":"status"}                              # operational/fault/mode/tool
    {"cmd":"drag_on","mode":"soft"}               # soft | floating
    {"cmd":"drag_off"}
    {"cmd":"joints"}
    {"cmd":"move_joints","q":[0,0,0,0,0,0,0],"speed":0.15}
    {"cmd":"quit"}

每个命令的应答都是单行 JSON，含 ``ok`` 字段；失败时另有 ``error``。
"""

from __future__ import annotations

import json
import os
import sys
import threading
import time
from pathlib import Path
from typing import List, Optional

# 公共 RDK 操作统一在 scripts/lib/rdk_helpers.py，避免与 rdk_switch_tool.py
# 各写一份导致行为漂移。
sys.path.insert(0, str(Path(__file__).resolve().parent / 'lib'))
from rdk_helpers import (  # noqa: E402
    clear_fault as rdk_clear_fault,
    ensure_idle,
    ensure_operational,
    rdk_whitelist_lists,
    resolve_tool_name,
    single_arm_groups,
    switch_tool,
    tool_name,
    tool_params,
)

_JSON_OUT = sys.stdout
sys.stdout = sys.stderr


def reply(obj: dict) -> None:
    _JSON_OUT.write(json.dumps(obj, ensure_ascii=False) + '\n')
    _JSON_OUT.flush()


def log(msg: str) -> None:
    print(f'[rdk_bridge] {msg}', file=sys.stderr, flush=True)


class SoftDrag:
    def __init__(self, robot, flexivrdk, lock: threading.Lock) -> None:
        self.robot = robot
        self.flexivrdk = flexivrdk
        self._lock = lock
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self.latest_q: List[float] = []
        self._groups = {}

    def start(self) -> None:
        mode = self.flexivrdk.Mode
        robot = self.robot
        with self._lock:
            robot.SwitchMode(mode.NRT_JOINT_IMPEDANCE)
            st = robot.states()
            self.latest_q = [float(v) for v in list(st.q)[:7]]
            try:
                info = robot.info()
                self._groups = dict(getattr(info, 'single_arm_groups', None) or {})
            except Exception as exc:
                log(f'info.single_arm_groups: {exc}')
                self._groups = {}
            if self._groups:
                for g in self._groups:
                    try:
                        nom = list(info.K_q_nom[g])
                        soft = [max(8.0, 0.08 * float(v)) for v in nom]
                        robot.SetJointImpedance(g, soft)
                        log(f'SetJointImpedance {g} K={soft}')
                    except Exception as exc:
                        log(f'SetJointImpedance 组 {g}: {exc}')
            else:
                n = max(7, len(self.latest_q))
                try:
                    robot.SetJointImpedance([35.0] * n, [0.7] * n)
                except Exception as exc:
                    log(f'SetJointImpedance flat（不分组）: {exc}')

        self._stop.clear()
        self._thread = threading.Thread(target=self._loop, name='soft-drag', daemon=True)
        self._thread.start()

    def _loop(self) -> None:
        robot = self.robot
        flexivrdk = self.flexivrdk
        while not self._stop.is_set():
            try:
                with self._lock:
                    if self._stop.is_set():
                        break
                    st = robot.states()
                    q = [float(v) for v in list(st.q)[:7]]
                    self.latest_q = q
                    if self._groups:
                        cmds = {}
                        for g in self._groups:
                            try:
                                gq = list(st[g].q)
                            except Exception:
                                gq = q
                            n = len(gq)
                            cmds[g] = flexivrdk.NrtJointPositionCmd(
                                gq, [0.0] * n, [1.5] * n, [2.0] * n
                            )
                        robot.SendJointPosition(cmds)
                    else:
                        n = len(q)
                        robot.SendJointPosition(q, [0.0] * n, [1.5] * n, [2.0] * n)
            except Exception as exc:
                log(f'soft 循环: {exc}')
                time.sleep(0.15)
            # 约 20 Hz —— 大部分时间让锁保持释放状态
            for _ in range(10):
                if self._stop.is_set():
                    return
                time.sleep(0.005)

    def stop(self) -> None:
        self._stop.set()
        th = self._thread
        if th is not None:
            th.join(timeout=4.0)
            if th.is_alive():
                log('join 之后 soft 循环仍然存活（无论如何都会切到 IDLE）')
            self._thread = None


def _ensure_operational(robot) -> Optional[str]:
    """保留旧名以兼容既有调用点；实现已统一到 rdk_helpers.ensure_operational。"""
    return ensure_operational(robot)


def _restore_tracking_impedance(robot, flexivrdk) -> None:
    """把关节阻抗从软拖拽拉回名义值，否则回放点位时臂跟不住指令。"""
    try:
        info = robot.info()
        groups = dict(getattr(info, 'single_arm_groups', None) or {})
    except Exception as exc:
        log(f'读名义阻抗失败: {exc}')
        groups = {}
    if groups:
        for group in groups:
            try:
                nominal = [float(v) for v in list(info.K_q_nom[group])]
                robot.SetJointImpedance(group, nominal)
                log(f'恢复阻抗 {group} K={nominal}')
            except Exception as exc:
                log(f'恢复阻抗 {group}: {exc}')
        return
    try:
        robot.SetJointImpedance([350.0] * 7, [7.0] * 7)
    except Exception as exc:
        log(f'恢复阻抗: {exc}')


def _send_joint_target(robot, flexivrdk, target: List[float], grouped: Optional[bool]) -> bool:
    """下发 7 关节目标。返回是否走了分组接口。"""
    if grouped is True:
        _send_grouped_joint_target(robot, flexivrdk, target)
        return True
    try:
        n = len(target)
        robot.SendJointPosition(target, [0.0] * n, [1.5] * n, [2.0] * n)
        return False
    except Exception as exc:
        if grouped is False:
            raise
        log(f'扁平关节指令失败，改分组: {exc}')
        _send_grouped_joint_target(robot, flexivrdk, target)
        return True


def _send_grouped_joint_target(robot, flexivrdk, target: List[float]) -> None:
    info = robot.info()
    groups = dict(getattr(info, 'single_arm_groups', None) or {})
    if not groups:
        raise RuntimeError('没有分组，扁平关节指令也失败了')
    states = robot.states()
    cmds = {}
    placed = False
    for group in groups:
        try:
            current = [float(v) for v in list(states[group].q)]
        except Exception:
            current = list(target)
        use = target if (not placed and len(current) == len(target)) else current
        if use is target:
            placed = True
        n = len(use)
        cmds[group] = flexivrdk.NrtJointPositionCmd(use, [0.0] * n, [1.5] * n, [2.0] * n)
    if not placed:
        raise RuntimeError('没有一组关节数和示教点一致，拒绝下发')
    robot.SendJointPosition(cmds)


def _idle_async(robot, lock: threading.Lock, flexivrdk, timeout: float = 4.0) -> bool:
    """切到 IDLE，同时避免把 stdin 循环永久阻塞住。"""
    done = threading.Event()

    def run() -> None:
        try:
            acquired = lock.acquire(timeout=2.0)
            try:
                if acquired:
                    try:
                        robot.Stop()
                    except Exception as exc:
                        log(f'Stop: {exc}')
                    try:
                        robot.SwitchMode(flexivrdk.Mode.IDLE)
                    except Exception as exc:
                        log(f'SwitchMode IDLE: {exc}')
            finally:
                if acquired:
                    lock.release()
        finally:
            done.set()

    threading.Thread(target=run, name='idle', daemon=True).start()
    ok = done.wait(timeout)
    if not ok:
        log('IDLE/Stop 超时 —— stdin 循环继续运行')
    return ok


def main() -> int:
    import flexivrdk

    robot = None
    drag_on = False
    soft: Optional[SoftDrag] = None
    rlock = threading.Lock()
    cached_q: List[float] = []

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except Exception as exc:  # noqa: BLE001
            reply({'ok': False, 'error': f'JSON 解析失败: {exc}'})
            continue

        cmd = str(msg.get('cmd', '')).strip()
        try:
            if cmd == 'ping':
                reply({'ok': True, 'py': sys.version.split()[0]})

            elif cmd == 'connect':
                sn = str(msg.get('sn') or os.environ.get('SITE_ROBOT_SN') or '')
                iface = str(msg.get('iface') or '')
                robot = None
                last_err = None
                attempts = []
                # 白名单必须是 IPv4。传 eth0 / enx… 会被 FastDDS 整段滤掉，
                # ping 通也会 Failed to connect。
                for listen in rdk_whitelist_lists(iface):
                    log('白名单: ' + ', '.join(listen))
                    attempts.extend([(sn, listen), (sn, listen, True)])
                attempts.append((sn,))
                for args in attempts:
                    try:
                        robot = flexivrdk.Robot(*args)
                        log(f'已连接 {sn} args={args!r}')
                        last_err = None
                        break
                    except TypeError:
                        continue
                    except Exception as exc:
                        last_err = exc
                        log(f'连接 {args!r}: {exc}')
                        robot = None
                if robot is None:
                    reply({'ok': False, 'error': str(last_err or '连接失败')})
                    continue
                err = _ensure_operational(robot)
                if err:
                    reply({'ok': False, 'error': err})
                    continue
                with rlock:
                    q = [float(v) for v in list(robot.states().q)[:7]]
                cached_q = q
                reply({'ok': True, 'sn': sn, 'q': q})

            elif cmd == 'drag_on':
                if robot is None:
                    reply({'ok': False, 'error': '尚未连接'})
                    continue
                err = _ensure_operational(robot)
                if err:
                    reply({'ok': False, 'error': err})
                    continue

                prefer = str(msg.get('mode', 'soft')).strip().lower()
                if prefer in ('floating', 'primitive', 'floatingjoint'):
                    name = str(msg.get('primitive', 'FloatingJoint')).replace('()', '')
                    try:
                        with rlock:
                            robot.SwitchMode(flexivrdk.Mode.NRT_PRIMITIVE_EXECUTION)
                            robot.ExecutePrimitive(name, {})
                        drag_on = True
                        reply({'ok': True, 'drag': True, 'mode': 'floating', 'primitive': name})
                        continue
                    except Exception as exc:
                        log(f'Floating 失败，退回 soft 模式: {exc}')

                try:
                    if soft is not None:
                        soft.stop()
                    soft = SoftDrag(robot, flexivrdk, rlock)
                    soft.start()
                    if soft.latest_q:
                        cached_q = list(soft.latest_q)
                    drag_on = True
                    log('软拖拽已开启')
                    reply({
                        'ok': True,
                        'drag': True,
                        'mode': 'soft',
                        'hint': '轻推机械臂即可，关节阻抗已调低。Enter=拍照（s=保持）。',
                    })
                except Exception as exc:
                    reply({'ok': False, 'error': f'软拖拽失败: {exc}'})

            elif cmd == 'drag_off':
                log('收到 drag_off 请求')
                if soft is not None:
                    if soft.latest_q:
                        cached_q = list(soft.latest_q)
                    soft.stop()
                    soft = None
                if robot is not None:
                    _idle_async(robot, rlock, flexivrdk, timeout=4.0)
                drag_on = False
                log('drag_off 已完成')
                reply({'ok': True, 'drag': False, 'q': cached_q})

            elif cmd == 'joints':
                q: Optional[List[float]] = None
                if soft is not None and len(soft.latest_q) >= 7:
                    q = list(soft.latest_q)
                elif len(cached_q) >= 7:
                    q = list(cached_q)
                if q is None or len(q) < 7:
                    reply({'ok': False, 'error': '没有关节缓存'})
                    continue
                cached_q = q
                reply({'ok': True, 'q': q, 'drag': drag_on})

            elif cmd == 'enable':
                if robot is None:
                    reply({'ok': False, 'error': '尚未连接'})
                    continue
                err = ensure_operational(robot)
                reply({
                    'ok': err is None,
                    'error': err,
                    'operational': bool(robot.operational()),
                })

            elif cmd == 'clear_fault':
                if robot is None:
                    reply({'ok': False, 'error': '尚未连接'})
                    continue
                err = rdk_clear_fault(robot)
                if err:
                    reply({'ok': False, 'error': err})
                    continue
                # 清故障后顺手使能，否则调用方还得再发一条 enable。
                err2 = ensure_operational(robot)
                reply({
                    'ok': err2 is None,
                    'error': err2,
                    'fault': bool(robot.fault()),
                    'operational': bool(robot.operational()),
                })

            elif cmd == 'switch_tool':
                if robot is None:
                    reply({'ok': False, 'error': '尚未连接'})
                    continue
                want = str(msg.get('name', '')).strip()

                # 切换工具要求 IDLE；这一步是同步的，因为切换本身不能被打断。
                err = ensure_idle(robot, flexivrdk.Mode, timeout_s=12.0)
                if err:
                    reply({'ok': False, 'error': err})
                    continue

                try:
                    with rlock:
                        tool = flexivrdk.Tool(robot)
                        groups = single_arm_groups(robot)
                        available = [str(n) for n in (tool.list() or [])]
                        anchor = groups[0] if groups else None

                        if not want:
                            # 纯查询：不切换，把当前工具与参数报回去。
                            params = tool_params(tool, anchor)
                            reply({
                                'ok': True,
                                'switched': False,
                                'current': tool_name(tool, anchor),
                                'available': available,
                                **params,
                            })
                            continue

                        target = resolve_tool_name(tool, want)
                        already = all(tool_name(tool, g) == target for g in groups)
                        if not already:
                            for g in groups:
                                switch_tool(tool, g, target)
                            time.sleep(0.3)

                        params = tool_params(tool, anchor)
                        reply({
                            'ok': True,
                            'switched': not already,
                            'already': already,
                            'current': tool_name(tool, anchor),
                            'target': target,
                            'available': available,
                            **params,
                        })
                except Exception as exc:  # noqa: BLE001
                    reply({'ok': False, 'error': str(exc)})

            elif cmd == 'status':
                if robot is None:
                    reply({'ok': False, 'connected': False, 'error': '尚未 connect'})
                    continue
                info = {'ok': True, 'connected': True, 'drag': drag_on}
                try:
                    info['operational'] = bool(robot.operational())
                    info['fault'] = bool(robot.fault())
                    info['mode'] = str(robot.mode())
                except Exception as exc:  # noqa: BLE001
                    info['error'] = str(exc)
                try:
                    with rlock:
                        tool = flexivrdk.Tool(robot)
                        groups = single_arm_groups(robot)
                        info['tool'] = tool_name(tool, groups[0] if groups else None)
                except Exception:
                    info['tool'] = ''
                reply(info)

            elif cmd == 'move_joints':
                if robot is None:
                    reply({'ok': False, 'error': '尚未连接'})
                    continue
                raw = msg.get('q')
                if not isinstance(raw, list) or len(raw) < 7:
                    reply({'ok': False, 'error': '需要 7 个关节角（弧度）'})
                    continue
                try:
                    target = [float(v) for v in raw[:7]]
                except (TypeError, ValueError):
                    reply({'ok': False, 'error': '关节角不是数字'})
                    continue
                if any(v != v or abs(v) > 12.6 for v in target):
                    reply({'ok': False, 'error': '关节角不是有限值，或超出约两圈'})
                    continue
                try:
                    speed = float(msg.get('speed', 0.15))
                except (TypeError, ValueError):
                    speed = 0.15
                speed = max(0.05, min(0.35, speed))
                if soft is not None:
                    soft.stop()
                    soft = None
                drag_on = False
                err = _ensure_operational(robot)
                if err:
                    reply({'ok': False, 'error': err})
                    continue
                try:
                    with rlock:
                        robot.SwitchMode(flexivrdk.Mode.NRT_JOINT_IMPEDANCE)
                        current = [float(v) for v in list(robot.states().q)[:7]]
                        _restore_tracking_impedance(robot, flexivrdk)
                    max_delta = max(abs(a - b) for a, b in zip(current, target))
                    step_s = 0.05
                    steps = max(1, int(max_delta / (speed * step_s)))
                    if steps * step_s > 150.0:
                        reply({
                            'ok': False,
                            'error': f'路程太长（{max_delta:.2f} rad），请先拖近一点再回放',
                        })
                        continue
                    grouped: Optional[bool] = None
                    for index in range(1, steps + 1):
                        alpha = index / steps
                        sample = [a + (b - a) * alpha for a, b in zip(current, target)]
                        with rlock:
                            grouped = _send_joint_target(robot, flexivrdk, sample, grouped)
                        time.sleep(step_s)
                    for _hold in range(8):
                        with rlock:
                            grouped = _send_joint_target(robot, flexivrdk, target, grouped)
                        time.sleep(step_s)
                    cached_q = list(target)
                    reply({'ok': True, 'q': target, 'from': current, 'steps': steps, 'speed': speed})
                except Exception as exc:
                    reply({'ok': False, 'error': f'回点失败: {exc}'})

            elif cmd == 'quit':
                if soft is not None:
                    soft.stop()
                    soft = None
                if robot is not None and drag_on:
                    _idle_async(robot, rlock, flexivrdk, timeout=3.0)
                reply({'ok': True})
                return 0

            else:
                reply({'ok': False, 'error': f'未知命令 {cmd}'})
        except Exception as exc:  # noqa: BLE001
            log(f'cmd={cmd} 出错: {exc}')
            reply({'ok': False, 'error': str(exc)})
    return 0


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