#!/usr/bin/env bash
# 把 Rizon4 移动到指定命名的位姿（默认：home）并强制停止。
#
# 关键：必须先杀掉 visual_servo.launch / servo_node，否则
# 它们会一直发 JointTrajectory，机械臂永远稳定不下来。
#
# 用法：
#   bash ~/ros2/scripts/go_home.sh
#   bash ~/ros2/scripts/go_home.sh --pose ready --speed 0.3

set -eo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
POSE="${POSE:-}"
SPEED="${SPEED:-}"
KEEP_SERVO=0

# shellcheck disable=SC1091
source "${ROOT}/scripts/lib/parse_arm_args.sh"
parse_arm_args "$@"
set -- "${PARSE_ARM_REST[@]}"

usage() {
  cat <<'EOF'
Move arm to a named pose (default: home), then hold with zero velocity.

Default: kill perception/servo/launch first (required to stop shaking).

Options:
  --pose NAME     home | ready | inspect | pos_left
  --speed FLOAT   velocity_scale (0,1]   default 0.25
  --arm 89|91     与 run_real 同一臂
  --keep-servo    do NOT kill servo (arm may keep shaking)
  -h, --help
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --pose) POSE="${2:?}"; shift 2 ;;
    --speed) SPEED="${2:?}"; shift 2 ;;
    --keep-servo) KEEP_SERVO=1; shift ;;
    --stop-track) shift ;; # 历史遗留
    -h|--help) usage; exit 0 ;;
    *) echo "Unknown: $1" >&2; usage >&2; exit 1 ;;
  esac
done

if ! command -v ros2 >/dev/null 2>&1 || [[ -z "${ROS_DISTRO:-}" ]]; then
  # shellcheck disable=SC1091
  source "${ROOT}/scripts/env.sh"
fi
# shellcheck disable=SC1091
source "${ROOT}/scripts/lib/load_site_config.sh"
POSE="${POSE:-${SITE_GO_HOME_POSE:-home}}"
SPEED="${SPEED:-${SITE_GO_HOME_SPEED:-0.25}}"

kill_streamers() {
  echo "=== killing visual-servo streamers ==="
  # 能先禁用控制就禁用（如果节点还活着），再动手 kill
  ros2 param set /servo_node enable_control false >/dev/null 2>&1 || true

  # 先杀 launch 父进程，再杀可执行文件（先 SIGTERM 再 SIGKILL）
  local patterns=(
    'visual_servo.launch.py'
    '/lib/arm_application/servo_node'
    '/lib/arm_application/perception_node'
  )
  local p
  for p in "${patterns[@]}"; do
    if pgrep -f "$p" >/dev/null 2>&1; then
      echo "  TERM $p"
      pkill -TERM -f "$p" 2>/dev/null || true
    fi
  done
  sleep 1.0
  for p in "${patterns[@]}"; do
    if pgrep -f "$p" >/dev/null 2>&1; then
      echo "  KILL $p"
      pkill -KILL -f "$p" 2>/dev/null || true
    fi
  done
  sleep 0.5

  if pgrep -f '/lib/arm_application/servo_node' >/dev/null 2>&1; then
    echo "ERROR: servo_node still alive — refuse to move (would keep shaking)" >&2
    pgrep -af '/lib/arm_application/servo_node' >&2 || true
    exit 2
  fi
  echo "  servo_node gone OK"
}

print_joints() {
  python3 <<'PY'
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import JointState
rclpy.init()
n = Node("go_home_joints")
box = [None]
n.create_subscription(JointState, "joint_states", lambda m: box.__setitem__(0, m), 10)
for _ in range(40):
    rclpy.spin_once(n, timeout_sec=0.05)
    if box[0] is not None:
        break
if box[0] is None:
    print("  (no joint_states)")
else:
    m = dict(zip(box[0].name, box[0].position))
    v = dict(zip(box[0].name, box[0].velocity)) if box[0].velocity else {}
    for i in range(1, 8):
        k = f"{__import__('os').environ.get('SITE_JOINT_PREFIX','Rizon4-063989_')}joint{i}"
        if k in m:
            print(f"  joint{i}: {m[k]:+.4f}  qd={v.get(k, 0.0):+.4f}")
n.destroy_node(); rclpy.shutdown()
PY
}

# 用控制器的 ACTION 做强力稳定（不走话题 —— 这样能可靠地匹配 JTC）
settle_hold_action() {
  local pose="$1"
  python3 - "$pose" <<'PY'
import sys, time
import rclpy
from rclpy.action import ActionClient
from rclpy.node import Node
from control_msgs.action import FollowJointTrajectory
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from sensor_msgs.msg import JointState

POSES = {
    "home": [0.0, -0.698, 0.0, 1.57, 0.0, 0.698, 0.0],
    "ready": [0.0, -0.698, 0.0, 1.57, 0.0, 0.698, 0.0],
    "inspect": [0.698, -0.698, 0.698, 1.57, 0.698, 0.698, 0.698],
    "pos_left": [-0.698, -0.698, 0.698, 1.57, 0.698, 0.698, 0.698],
}
name = sys.argv[1]
target = POSES.get(name)
if target is None:
    raise SystemExit(f"unknown pose {name}")
import os
PREFIX = os.environ.get("SITE_JOINT_PREFIX", "Rizon4-063989_")
JOINTS = [f"{PREFIX}joint{i}" for i in range(1, 8)]

rclpy.init()
n = Node("go_home_hold_action")
js_box = [None]
n.create_subscription(JointState, "joint_states", lambda m: js_box.__setitem__(0, m), 10)
cli = ActionClient(n, FollowJointTrajectory, "rizon_arm_controller/follow_joint_trajectory")
if not cli.wait_for_server(timeout_sec=5.0):
    print("  WARN: JTC action not available — skip action hold")
    n.destroy_node(); rclpy.shutdown(); raise SystemExit(0)

def spin_until(pred, timeout=3.0):
    t0 = time.monotonic()
    while time.monotonic() - t0 < timeout:
        rclpy.spin_once(n, timeout_sec=0.05)
        if pred():
            return True
    return False

spin_until(lambda: js_box[0] is not None, 3.0)

# 两点轨迹：当前位姿 -> 目标位姿（或保持在目标位姿），两点速度都为零
q_now = list(target)
if js_box[0] is not None:
    m = dict(zip(js_box[0].name, js_box[0].position))
    try:
        q_now = [float(m[j]) for j in JOINTS]
    except KeyError:
        pass

def send_hold(q, T=0.40):
    traj = JointTrajectory()
    traj.joint_names = list(JOINTS)
    # 起点
    p0 = JointTrajectoryPoint()
    p0.positions = [float(v) for v in q]
    p0.velocities = [0.0] * 7
    p0.accelerations = [0.0] * 7
    p0.time_from_start.sec = 0
    p0.time_from_start.nanosec = 50_000_000
    # 终点是同一姿势，速度仍为零（钉住）
    p1 = JointTrajectoryPoint()
    p1.positions = [float(v) for v in q]
    p1.velocities = [0.0] * 7
    p1.accelerations = [0.0] * 7
    ns = int(round(T * 1e9))
    p1.time_from_start.sec = ns // 1_000_000_000
    p1.time_from_start.nanosec = ns % 1_000_000_000
    traj.points = [p0, p1]
    goal = FollowJointTrajectory.Goal()
    goal.trajectory = traj
    fut = cli.send_goal_async(goal)
    spin_until(lambda: fut.done(), 3.0)
    if not fut.done():
        print("  hold goal send timeout")
        return False
    gh = fut.result()
    if gh is None or not gh.accepted:
        print("  hold goal rejected")
        return False
    rfut = gh.get_result_async()
    spin_until(lambda: rfut.done(), T + 3.0)
    return True

# 1) 先朝命名位姿快速靠拢，末端速度为零
print(f"  action hold → {name} (zero vel)")
send_hold(target, T=0.60)
time.sleep(0.4)
# 2–4) 反复把实测位姿重新钉住（让 JTC 的开环指令与编码器读数对齐）
for i in range(3):
    spin_until(lambda: js_box[0] is not None, 1.0)
    if js_box[0] is None:
        break
    m = dict(zip(js_box[0].name, js_box[0].position))
    q_pin = [float(m[j]) for j in JOINTS]
    print(f"  action re-pin #{i+1} @ measured joints (zero vel)")
    send_hold(q_pin, T=0.45)
    time.sleep(0.25)

# 等到 |qd| 足够小
print("  waiting for joints to stop ...")
t0 = time.monotonic()
while time.monotonic() - t0 < 5.0:
    rclpy.spin_once(n, timeout_sec=0.05)
    if js_box[0] is None or not js_box[0].velocity:
        continue
    vmap = dict(zip(js_box[0].name, js_box[0].velocity))
    vmax = max(abs(float(vmap.get(j, 0.0))) for j in JOINTS)
    if vmax < 0.015:
        print(f"  settled |qd|_max={vmax:.4f} rad/s")
        break
else:
    print("  WARN: still some residual velocity")

n.destroy_node()
rclpy.shutdown()
PY
}

if [[ "${KEEP_SERVO}" -eq 0 ]]; then
  kill_streamers
else
  echo "WARNING: --keep-servo — shaking likely"
fi

if ! ros2 service list 2>/dev/null | grep -q '/set_named_pose'; then
  echo "ERROR: /set_named_pose missing — start run_real.sh" >&2
  exit 1
fi

echo "=== joints before ==="
print_joints

echo "=== move: ${POSE} speed=${SPEED} ==="
RESP="$(
  ros2 service call /set_named_pose arm_msgs/srv/SetNamedPose \
    "{pose_name: '${POSE}', velocity_scale: ${SPEED}, acceleration_scale: 0.0}" 2>&1
)"
echo "${RESP}"
if ! echo "${RESP}" | grep -q 'success=True'; then
  echo "ERROR: move failed" >&2
  exit 1
fi

WAIT_SEC="$(python3 -c "print(max(2.5, 1.0 / max(float('${SPEED}'), 0.05) + 2.0))")"
echo "Waiting ${WAIT_SEC}s ..."
sleep "${WAIT_SEC}"

echo "=== settle (action, vel=0) ==="
settle_hold_action "${POSE}"

# 最后确认：servo 必须仍然是死的
if pgrep -f '/lib/arm_application/servo_node' >/dev/null 2>&1; then
  echo "WARNING: servo_node came back — killing again"
  kill_streamers
  settle_hold_action "${POSE}"
fi

echo "=== joints after ==="
print_joints
echo "Done. Arm should be still."
echo "Tracking later: bash ~/ros2/scripts/run_visual_servo_rt.sh --control"