#!/usr/bin/env bash
# MuJoCo O30i 左手单窗仿真
# URDF: exports/O30_V9_left_right/left/linkerhand_o30_left_v9.urdf (V10 right mirrored; existing asset path retained)
# 订阅 /cb_left_hand_control_cmd（20 路 SDK 0～255，按 URDF 限位线性换算）
# 发布 /o30/left/joint_states（20 路弧度）
#
# 用法:
#   ./hands/O30/scripts/run_o30_left_mujoco_plan.sh
#   ./hands/O30/scripts/run_o30_left_mujoco_plan.sh --pd-drive
set -eo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
cd "$ROOT"

URDF_SRC="${O30_LEFT_URDF:-$ROOT/assets/exports/O30_V10_left_right/left/linkerhand_o30_left_v10.urdf}"
MESH_SRC="${O30_LEFT_MESHDIR:-$(dirname "$URDF_SRC")/meshes}"
PKG_SRC="$ROOT/src/linkerhand-sim/linker_hand_mujoco_ros2/linker_hand_mujoco_ros2"
O30_DIR="$PKG_SRC/urdf/O30/linker_hand_o30_left"
XML="$O30_DIR/linker_hand_o30_left.xml"
URDF_PKG="$O30_DIR/linkerhand_o30_left.urdf"

if [[ ! -f "$URDF_SRC" ]]; then
  echo "[error] missing O30 URDF: $URDF_SRC" >&2
  exit 1
fi
if [[ ! -d "$MESH_SRC" ]]; then
  echo "[error] missing O30 meshes: $MESH_SRC" >&2
  exit 1
fi

set +u
# shellcheck disable=SC1091
source /opt/ros/jazzy/setup.bash
if [[ -f "$ROOT/install/setup.bash" ]]; then
  # shellcheck disable=SC1091
  source "$ROOT/install/setup.bash"
fi
set -u

VENV_SITE="$ROOT/venv/lib/python3.12/site-packages"
if [[ -d "$VENV_SITE" ]]; then
  export PYTHONPATH="$VENV_SITE${PYTHONPATH:+:$PYTHONPATH}"
fi
export PYTHONPATH="$ROOT/src/linkerhand-sim/linker_hand_mujoco_ros2${PYTHONPATH:+:$PYTHONPATH}"

echo "[o30-left] sync MJCF from $URDF_SRC"
python3 - "$URDF_SRC" "$MESH_SRC" "$O30_DIR" "$XML" "$URDF_PKG" <<'PY'
import re
import shutil
import sys
from pathlib import Path

import mujoco

urdf_src = Path(sys.argv[1])
mesh_src = Path(sys.argv[2])
out_dir = Path(sys.argv[3])
xml_path = Path(sys.argv[4])
urdf_pkg = Path(sys.argv[5])

out_dir.mkdir(parents=True, exist_ok=True)
mesh_out = out_dir / "meshes"
mesh_out.mkdir(exist_ok=True)
for f in mesh_src.iterdir():
    if f.is_file():
        shutil.copy2(f, mesh_out / f.name)

text = urdf_src.read_text()
text = re.sub(r"package://[^/]+/meshes/", "meshes/", text)
urdf_pkg.write_text(text)


def attrs_of(chunk: str) -> dict[str, str]:
    return dict(re.findall(r'(\w+)="([^"]*)"', chunk))


def parse_links(urdf: str) -> dict[str, dict]:
    links: dict[str, dict] = {}
    for m in re.finditer(r'<link\s+name="([^"]+)"[^>]*>(.*?)</link>', urdf, re.S):
        name, body = m.group(1), m.group(2)
        inertial = re.search(r"<inertial>(.*?)</inertial>", body, re.S)
        pos = "0 0 0"
        mass = "0.01"
        fullinertia = "1e-6 1e-6 1e-6 0 0 0"
        if inertial:
            ib = inertial.group(1)
            origin = re.search(r"<origin\b([^/]*)/>", ib)
            mass_m = re.search(r'<mass\b([^/]*)/>', ib)
            inertia = re.search(r"<inertia\b([^/]*)/>", ib)
            if origin:
                pos = attrs_of(origin.group(1)).get("xyz", pos)
            if mass_m:
                mass = attrs_of(mass_m.group(1)).get("value", mass)
            if inertia:
                ia = attrs_of(inertia.group(1))
                fullinertia = " ".join(
                    ia.get(k, "0") for k in ("ixx", "iyy", "izz", "ixy", "ixz", "iyz")
                )
        color = "0.7 0.7 0.7 1"
        vis = re.search(r"<visual>(.*?)</visual>", body, re.S)
        if vis:
            col = re.search(r'rgba="([^"]+)"', vis.group(1))
            if col:
                color = col.group(1)
        mesh = None
        mesh_m = re.search(r'<mesh\b([^>]*)/>', body)
        if mesh_m:
            fn = attrs_of(mesh_m.group(1)).get("filename", "")
            mesh = Path(fn).name
        links[name] = {
            "pos": pos,
            "mass": mass,
            "fullinertia": fullinertia,
            "color": color,
            "mesh": mesh,
        }
    return links


def parse_joints(urdf: str) -> list[dict]:
    joints = []
    for m in re.finditer(r"<joint\b([^>]*)>(.*?)</joint>", urdf, re.S):
        attrs, body = m.group(1), m.group(2)
        name = re.search(r'name="([^"]+)"', attrs)
        jtype = re.search(r'type="([^"]+)"', attrs)
        if not name or not jtype or jtype.group(1) != "revolute":
            continue
        origin = re.search(r"<origin\b([^/]*)/>", body)
        axis = re.search(r"<axis\b([^/]*)/>", body)
        lim = re.search(r"<limit\b([^/]*)/>", body)
        parent = re.search(r'<parent[^>]*link="([^"]+)"', body)
        child = re.search(r'<child[^>]*link="([^"]+)"', body)
        oa = attrs_of(origin.group(1)) if origin else {}
        la = attrs_of(lim.group(1)) if lim else {}
        joints.append(
            {
                "name": name.group(1),
                "xyz": oa.get("xyz", "0 0 0"),
                "rpy": oa.get("rpy", "0 0 0"),
                "axis": attrs_of(axis.group(1)).get("xyz", "0 0 1") if axis else "0 0 1",
                "lo": float(la.get("lower", "0")),
                "hi": float(la.get("upper", "0")),
                "parent": parent.group(1) if parent else "",
                "child": child.group(1) if child else "",
            }
        )
    return joints


def rpy_to_quat(rpy: str) -> str | None:
    import math

    r, p, y = [float(x) for x in rpy.split()]
    if abs(r) + abs(p) + abs(y) < 1e-12:
        return None
    cr, sr = math.cos(r / 2), math.sin(r / 2)
    cp, sp = math.cos(p / 2), math.sin(p / 2)
    cy, sy = math.cos(y / 2), math.sin(y / 2)
    w = cr * cp * cy + sr * sp * sy
    x = sr * cp * cy - cr * sp * sy
    yq = cr * sp * cy + sr * cp * sy
    z = cr * cp * sy - sr * sp * cy
    return f"{w:.16g} {x:.16g} {yq:.16g} {z:.16g}"


links = parse_links(text)
joints = parse_joints(text)
children = {j["parent"]: [] for j in joints}
for j in joints:
    children.setdefault(j["parent"], []).append(j)
# keep URDF order
order = [j["name"] for j in joints]
by_parent: dict[str, list[dict]] = {}
for j in joints:
    by_parent.setdefault(j["parent"], []).append(j)

root = "hand_base_link"
meshes = sorted({info["mesh"] for info in links.values() if info["mesh"]})

lines = [
    '<mujoco model="linkerhand_o30_left">',
    '  <compiler angle="radian" meshdir="meshes"/>',
    '  <option timestep="0.002" gravity="0 0 -9.81" integrator="implicitfast"/>',
    '  <statistic center="0.01 0 0.08" extent="0.16"/>',
    "  <default>",
    '    <joint damping="0.15" armature="0.01"/>',
    '    <position kp="80" kv="8" forcelimited="true" forcerange="-8 8"/>',
    '    <default class="visual">',
    '      <geom type="mesh" contype="0" conaffinity="0" group="2"/>',
    "    </default>",
    '    <default class="collision">',
    '      <geom type="mesh" contype="1" conaffinity="1" condim="3" group="3"',
    '            rgba="0 1 0 0.15" friction="1 0.5 0.01"',
    '            solref="0.004 1" solimp="0.95 0.99 0.001" margin="0"/>',
    "    </default>",
    "  </default>",
    "  <asset>",
]
# Respect the asset scale: V10 STL vertices are already mirrored.
import xml.etree.ElementTree as ET
mesh_scales = {}
for element in ET.fromstring(text).findall('.//mesh'):
    filename = Path(element.get('filename')).name
    scale = element.get('scale', '1 1 1')
    if filename in mesh_scales and mesh_scales[filename] != scale:
        raise ValueError(f'Conflicting mesh scales: {filename}')
    mesh_scales[filename] = scale
for mesh in meshes:
    scale = mesh_scales[Path(mesh).name]
    lines.append(f'    <mesh name="{Path(mesh).stem}" file="{mesh}" scale="{scale}"/>')
lines.append("  </asset>")
lines.append("  <worldbody>")


def emit_body(link: str, indent: int, joint: dict | None) -> None:
    pad = "  " * indent
    info = links[link]
    if joint is None:
        lines.append(f'{pad}<body name="{link}">')
    else:
        quat = rpy_to_quat(joint["rpy"])
        pose = f'pos="{joint["xyz"]}"'
        if quat:
            pose += f' quat="{quat}"'
        lines.append(f'{pad}<body name="{link}" {pose}>')
    lines.append(
        f'{pad}  <inertial pos="{info["pos"]}" mass="{info["mass"]}" '
        f'fullinertia="{info["fullinertia"]}"/>'
    )
    if joint is not None:
        lines.append(
            f'{pad}  <joint name="{joint["name"]}" pos="0 0 0" axis="{joint["axis"]}" '
            f'range="{joint["lo"]:.16g} {joint["hi"]:.16g}"/>'
        )
    mesh = Path(info["mesh"]).stem if info["mesh"] else None
    if mesh:
        lines.append(
            f'{pad}  <geom class="visual" rgba="{info["color"]}" mesh="{mesh}"/>'
        )
        lines.append(f'{pad}  <geom class="collision" mesh="{mesh}"/>')
    for child_j in by_parent.get(link, []):
        emit_body(child_j["child"], indent + 1, child_j)
    lines.append(f"{pad}</body>")


emit_body(root, 2, None)
lines.append("  </worldbody>")
lines.append("  <contact>")
for j in joints:
    lines.append(f'    <exclude body1="{j["parent"]}" body2="{j["child"]}"/>')
# 张开位掌心与近节 mesh 贴合，否则 PD 会把手指顶开
for other in (
    "thumb_metacarpals_base2",
    "index_proximal",
    "middle_proximal",
    "ring_proximal",
    "pinky_proximal",
):
    lines.append(f'    <exclude body1="hand_base_link" body2="{other}"/>')
lines.append("  </contact>")
lines.append("  <actuator>")
for j in joints:
    lines.append(
        f'    <position name="{j["name"]}_pos" joint="{j["name"]}" '
        f'ctrlrange="{j["lo"]:.16g} {j["hi"]:.16g}"/>'
    )
lines.append("  </actuator>")
lines.append("</mujoco>")
xml_path.write_text("\n".join(lines) + "\n")
model = mujoco.MjModel.from_xml_path(str(xml_path))
print(f"[o30-left] wrote {xml_path} njnt={model.njnt} nbody={model.nbody}")
PY

if [[ "${O30_LEFT_PREPARE_ONLY:-0}" == "1" ]]; then
  exit 0
fi

echo "[o30-left] urdf=$URDF_PKG"
echo "[o30-left] xml=$XML"
echo "[o30-left] cmd-topic=/cb_left_hand_control_cmd (20 SDK values, 0..255)"
echo "[o30-left] state-topic=/o30/left/joint_states (20 rad)"
echo "[o30-left] SDK order: thumb_roll,thumb_yaw,index_yaw,middle_yaw,ring_yaw,pinky_yaw, thumb_root,index_root,middle_root,ring_root,pinky_root, index_pip,middle_pip,ring_pip,pinky_pip, thumb_tip,index_tip,middle_tip,ring_tip,pinky_tip"
echo "[o30-left] default mode=kinematic; 加 --pd-drive 看接触"

exec python3 -m linker_hand_mujoco_ros2.o30_left_plan_bridge \
  --xml "$XML" \
  --node-name o30_mujoco_plan_bridge_left \
  --window-title "O30i left V10 mirrored" \
  --cmd-topic /cb_left_hand_control_cmd \
  --sim-state-topic /o30/left/joint_states \
  --no-legacy-sdk \
  --mapping g20_urdf \
  "$@"