"""把双臂工作站 + 操作台 + 测试物体导出成浏览器能直接跑的自包含 MJCF。

  ~/unitree_rl_mjlab/.venv/bin/python tools/export_manip_scene.py

产出 scenes/<name>/:
  model.xml   自包含 MJCF(网格内联 + 视觉网格抽稀,不依赖任何外部文件)
  meta.json   给前端用的索引:臂/手的执行器、工具 site、指尖 geom、
              物体和桌面的几何参数、就绪位形

这里没有策略,是纯手动操作的沙盒 —— 前端用 IK 驱动腕部、用一个握力滑条驱动手指。

网格处理照搬 unitree_rl_mjlab/scripts/export_web_dog.py 里踩出来的三条:
· 写回顶点前要用 mesh_pos/mesh_quat 变换回原坐标系,否则重新编译会再对齐一次,
  几何整体偏移;
· MJCF 的 <mesh normal> 是**每顶点**一条,不是每面;
· spec 里源文件的法线不受 usernormal 覆盖,只能删掉网格重建。
"""

import json
from dataclasses import dataclass
from pathlib import Path

import mujoco
import numpy as np
import tyro

HERE = Path(__file__).resolve().parent.parent
WS_ROOT = (Path.home() / "linker-sim/packages/linker-robot-assets/src"
           / "linker_robot_assets/assets/workstations")


@dataclass
class Cfg:
  workstation: str = "p7_i1_l6_bimanual"
  """linker-robot-assets 里的工作站名。"""
  name: str = "p7_bimanual"
  """导出目录名。"""
  table_top: float = 0.75
  """操作台高度 [m]。"""
  table_x: float = 0.42
  """操作台中心到机器人的距离 [m]。"""
  decimate: int = 32
  """视觉网格的聚类抽稀格数;0 = 不抽稀。"""
  gripper: str = "so101"
  """末端:l6=原装五指手;so101=换成 SO-101 二指夹爪(仅 rig=p7 时有效)。"""
  rig: str = "so101_pair"
  """机器人:p7=linker 双臂工作站;so101_pair=桌面上并排两台 SO-101。"""


# ---------------------------------------------------------------- 网格处理

def decimate(v, f, cells=48):
  """顶点聚类抽稀:包围盒切 cells^3 的格子,每格顶点并成一个。"""
  lo, hi = v.min(0), v.max(0)
  size = float((hi - lo).max())
  if size <= 0:
    return v, f
  key = np.floor((v - lo) / (size / cells)).astype(np.int64)
  uniq, inv = np.unique(key, axis=0, return_inverse=True)
  acc = np.zeros((len(uniq), 3))
  cnt = np.zeros(len(uniq))
  np.add.at(acc, inv, v)
  np.add.at(cnt, inv, 1)
  f2 = inv[f]
  keep = (f2[:, 0] != f2[:, 1]) & (f2[:, 1] != f2[:, 2]) & (f2[:, 0] != f2[:, 2])
  f2 = f2[keep]
  if len(f2) < 4 or len(uniq) < 4:
    return v, f
  used, f3 = np.unique(f2, return_inverse=True)
  v2, f3 = (acc / cnt[:, None])[used], f3.reshape(-1, 3)
  # 抽稀把薄零件压成一张片时,MuJoCo 会拒收("mesh volume is too small")。
  # 退化了就原样返回,少省几百 KB 而已。
  ext = v2.max(0) - v2.min(0)
  if len(v2) < 20 or float(ext.min()) < 1e-3:
    return v, f
  return v2, f3


def convex_hull(v):
  """取凸包 —— MuJoCo 做网格碰撞时用的本来就是凸包,所以这一步零精度损失,
  但顶点数能从几万降到几百。"""
  from scipy.spatial import ConvexHull

  h = ConvexHull(v)
  used, faces = np.unique(h.simplices, return_inverse=True)
  return v[used], faces.reshape(-1, 3)


def inline_meshes(spec: mujoco.MjSpec, m: mujoco.MjModel, cells: int) -> None:
  """网格内联进 XML。

  这套资产的视觉几何和碰撞几何**引用同一批网格**(41 个全是两用),
  直接抽稀会改碰撞形状,不抽稀又有 49 万顶点 / 25MB。做法是拆成两份:
  · 碰撞用凸包(MuJoCo 反正只用凸包,行为完全不变)
  · 视觉用抽稀版 `<name>_vis`,再把纯视觉 geom 的 meshname 改过去
  """
  by_name = {m.mesh(i).name: i for i in range(m.nmesh)}
  collide, visual = set(), set()
  for g in range(m.ngeom):
    if m.geom_type[g] != mujoco.mjtGeom.mjGEOM_MESH:
      continue
    (collide if (m.geom_contype[g] or m.geom_conaffinity[g]) else visual).add(
      int(m.geom_dataid[g]))

  todo = []
  for mesh in list(spec.meshes):
    i = by_name.get(mesh.name)
    if i is None:
      continue
    va, vn = m.mesh_vertadr[i], m.mesh_vertnum[i]
    fa, fn = m.mesh_faceadr[i], m.mesh_facenum[i]
    v = m.mesh_vert[va:va + vn].astype(np.float64)
    rot = np.zeros(9)
    mujoco.mju_quat2Mat(rot, m.mesh_quat[i])
    v = v @ rot.reshape(3, 3).T + m.mesh_pos[i]       # 变换回原坐标系
    f = m.mesh_face[fa:fa + fn].astype(np.int64)
    todo.append((mesh.name, i, v, f))

  for nm, _, _, _ in todo:
    spec.delete(spec.mesh(nm))                        # 删掉重建,甩掉源文件的法线

  def emit(name, v, f):
    nu = spec.add_mesh()
    nu.name = name
    nu.uservert = np.round(v, 4).reshape(-1).tolist()
    nu.userface = f.astype(np.int32).reshape(-1).tolist()

  n_v0 = n_v1 = 0
  for nm, i, v, f in todo:
    n_v0 += len(v)
    if i in collide:
      try:
        cv, cf = convex_hull(v)
      except Exception:                               # noqa: BLE001
        cv, cf = decimate(v, f, 64)
      emit(nm, cv, cf)
      n_v1 += len(cv)
    else:
      emit(nm, *((decimate(v, f, cells)) if cells else (v, f)))
      n_v1 += len(spec.mesh(nm).uservert) // 3
    if i in visual and i in collide:
      vv, vf = decimate(v, f, cells) if cells else (v, f)
      emit(f"{nm}_vis", vv, vf)
      n_v1 += len(vv)

  # 纯视觉的 geom 指向抽稀版
  for b in spec.bodies:
    for g in b.geoms:
      if (g.contype == 0 and g.conaffinity == 0 and g.meshname
          and by_name.get(g.meshname) in visual
          and by_name.get(g.meshname) in collide):
        g.meshname = f"{g.meshname}_vis"
  print(f"  网格顶点 {n_v0} → {n_v1}")


# ---------------------------------------------------------------- 场景搭建

# DexGraspBench example 数据集里的五个物体(灵巧手抓取那套方案用的就是这五个)。
# 目录名很长,这里给一组短名和中文标签。
DGB_ROOT = Path.home() / "DexGraspBench/assets/example_object"
OBJECTS = [
  # 短名,     中文,     目录,                                              抓握截面 m, 质量 kg, rgba
  # 抓握截面按手来定:L6 手张开 13.4cm、**合拢只剩 4cm**,截面超过 4~5cm 就夹不住
  # (6.6cm 的罐子实测怎么放都会被挤出去)。这里统一缩到 4cm 上下。
  ("can",      "易拉罐", "core_can_be67418a10003cc9eae3efbc9dbeea",         0.040, 0.15, (0.85, 0.30, 0.25, 1)),
  ("bottle",   "瓶子",   "core_bottle_523cddb320608c09a37f3fc191551700",    0.038, 0.12, (0.30, 0.65, 0.85, 1)),
  ("jar",      "罐子",   "core_jar_58c9d6575d62fbaf12bc68f36f3bdd45",       0.042, 0.16, (0.55, 0.75, 0.35, 1)),
  ("piano",    "玩具琴", "sem_Piano_438f64dd9ff80bfddd5efa4c0e7932e",       0.045, 0.12, (0.85, 0.70, 0.25, 1)),
  ("toy",      "玩偶",   "sem_ToyFigure_a3087c9105fe0878d18a6b3b9816ca14",  0.040, 0.10, (0.70, 0.45, 0.80, 1)),
]


def load_obj(path):
  """极简 OBJ 读取:只要顶点和面(多边形按扇形三角化)。"""
  vs, fs = [], []
  for ln in Path(path).read_text().splitlines():
    if ln.startswith("v "):
      vs.append([float(x) for x in ln.split()[1:4]])
    elif ln.startswith("f "):
      idx = [int(t.split("/")[0]) - 1 for t in ln.split()[1:]]
      for k in range(1, len(idx) - 1):
        fs.append([idx[0], idx[k], idx[k + 1]])
  return np.array(vs, float), np.array(fs, np.int64)


def add_objects(spec, wb, top_z, ox, ys):
  """把五个测试物体摆成一排。ox 是它们离机器人的距离,ys 是横向位置。"""
  objs = []
  for (nm, label, folder, width, mass, rgba), oy in zip(OBJECTS, ys):
    root = DGB_ROOT / folder
    vv, vf = load_obj(root / "mesh" / "simplified.obj")
    ext = vv.max(0) - vv.min(0)
    # 按"抓握截面"定缩放:取次长边缩到 width,长边就是高度
    scale = width / float(sorted(ext)[1])
    vv = (vv - (vv.max(0) + vv.min(0)) / 2) * scale       # 顺便把几何中心挪到原点

    # 立起来:原始网格的长轴是 y,转 90° 让它朝 z
    rot = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]], float)
    vv = vv @ rot.T
    half_h = float((vv[:, 2].max() - vv[:, 2].min()) / 2)

    mv = spec.add_mesh()
    mv.name = f"obj_{nm}_vis"
    mv.uservert = np.round(vv, 5).reshape(-1).tolist()
    mv.userface = vf.astype(np.int32).reshape(-1).tolist()

    b = wb.add_body(name=f"obj_{nm}", pos=[ox, oy, top_z + half_h + 0.002])
    b.add_freejoint()
    # 视觉网格背质量(凸块按密度算出来的质量是凸分解伪影,重叠处被重复累加,
    # 实测能虚高好几倍 —— 直接显式给质量)
    b.add_geom(name=f"obj_{nm}", type=mujoco.mjtGeom.mjGEOM_MESH, meshname=mv.name,
               rgba=list(rgba), mass=mass, contype=0, conaffinity=0, group=2)

    pieces = sorted((root / "urdf" / "meshes").glob("convex_piece_*.obj"))
    for k, pf in enumerate(pieces):
      cv, cf = load_obj(pf)
      cv = ((cv - (vv.max(0) + vv.min(0)) * 0) * scale)     # 与视觉网格同一缩放
      cv = cv - (load_obj(root / "mesh" / "simplified.obj")[0].max(0)
                 + load_obj(root / "mesh" / "simplified.obj")[0].min(0)) / 2 * scale
      cv = cv @ rot.T
      mc = spec.add_mesh()
      mc.name = f"obj_{nm}_col{k}"
      mc.uservert = np.round(cv, 5).reshape(-1).tolist()
      mc.userface = cf.astype(np.int32).reshape(-1).tolist()
      b.add_geom(name=f"obj_{nm}_c{k}", type=mujoco.mjtGeom.mjGEOM_MESH,
                 meshname=mc.name, density=0, contype=1, conaffinity=1, condim=4,
                 friction=[1.2, 0.01, 0.0005], solref=[0.008, 1.0],
                 group=3, rgba=[1, 0, 0, 0.25])

    objs.append({"name": nm, "label": label, "body": f"obj_{nm}",
                 "geom": f"obj_{nm}", "kind": "mesh", "mass": mass,
                 "half_h": half_h, "pieces": len(pieces),
                 "home": [ox, oy, top_z + half_h + 0.002]})
  return objs


def add_scene(spec: mujoco.MjSpec, cfg: Cfg):
  """地面、灯光、操作台、三个测试物体。"""
  wb = spec.worldbody
  tex = spec.add_texture(
    name="grid", type=mujoco.mjtTexture.mjTEXTURE_2D,
    builtin=mujoco.mjtBuiltin.mjBUILTIN_CHECKER,
    rgb1=[0.17, 0.20, 0.25], rgb2=[0.12, 0.15, 0.19], width=300, height=300)
  del tex
  mat = spec.add_material(name="grid", texrepeat=[8, 8], reflectance=0.05)
  mat.textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] = "grid"
  wb.add_geom(name="floor", type=mujoco.mjtGeom.mjGEOM_PLANE, size=[0, 0, 0.05],
              material="grid", contype=1, conaffinity=1, condim=3)
  wb.add_light(pos=[0.4, 0, 3.0], dir=[0, 0, -1],
               type=mujoco.mjtLightType.mjLIGHT_DIRECTIONAL)

  # 操作台:台面 + 四条腿(纯装饰,只有台面参与碰撞)
  half = (0.28, 0.42, 0.02)
  top_z = cfg.table_top
  t = wb.add_body(name="table", pos=[cfg.table_x, 0, 0])
  t.add_geom(name="table_top", type=mujoco.mjtGeom.mjGEOM_BOX, size=list(half),
             pos=[0, 0, top_z - half[2]], rgba=[0.55, 0.42, 0.30, 1],
             contype=1, conaffinity=1, condim=3, friction=[1.0, 0.005, 0.0001])
  # 台面还要一份纯视觉的副本:渲染器的"视觉"过滤会把有视觉兄弟的碰撞几何藏起来
  # (桌腿是纯视觉的),不补这一份台面就画不出来,物体看着像飘在空中。
  t.add_geom(name="table_top_vis", type=mujoco.mjtGeom.mjGEOM_BOX,
             size=[half[0] * 1.001, half[1] * 1.001, half[2] * 0.999],
             pos=[0, 0, top_z - half[2]], rgba=[0.55, 0.42, 0.30, 1],
             contype=0, conaffinity=0, mass=0)
  for sx in (-1, 1):
    for sy in (-1, 1):
      t.add_geom(name=f"table_leg_{sx}_{sy}", type=mujoco.mjtGeom.mjGEOM_BOX,
                 size=[0.02, 0.02, (top_z - 2 * half[2]) / 2],
                 pos=[sx * (half[0] - 0.03), sy * (half[1] - 0.03),
                      (top_z - 2 * half[2]) / 2],
                 rgba=[0.38, 0.29, 0.21, 1], contype=0, conaffinity=0)

  objs = add_objects(spec, wb, top_z, cfg.table_x - 0.02, ys=[-0.26,-0.13,0.0,0.13,0.26])
  return {"table": {"x": cfg.table_x, "top": top_z, "half": list(half)},
          "objects": objs}


# ---------------------------------------------------------------- 换末端

SO101_XML = (Path.home() / "unitree_rl_mjlab/src/assets/robots/unitree_go2_so101"
             / "xmls/so101/so101.xml")
SO101_KP, SO101_KV, SO101_FRC = 998.22, 2.731, 2.94


def swap_gripper(spec: mujoco.MjSpec):
  """把 L6 五指手换成 SO-101 的二指夹爪。

  五指手在这套场景里抓不住东西:张开 13.4cm、合拢只剩 4cm,四指从张开位置
  扫向掌心的路径正好经过物体,合拢那一下必把物体挤飞(试过 18 组瞄准位置、
  三种掌心朝向、拇指先合的分阶段闭合,都不行)。二指夹爪是已经在 Go2 上
  验证过的方案,直接把它的 `gripper` 子树挂到腕部安装点。
  """
  info = {}
  for side in ("left", "right"):
    # 每侧读一份:attach_body 会把源 spec 里的 body 取走,复用同一份第二次就是 None
    grip_src = mujoco.MjSpec.from_file(str(SO101_XML))
    for act in list(grip_src.actuators):
      grip_src.delete(act)                   # 执行器这边自己重建
    pre = "l" if side == "left" else "r"
    hand = spec.body(f"hand_{side}_{pre}h_hand_base_link")
    parent, pos, quat = hand.parent, list(hand.pos), list(hand.quat)
    # 先删执行器再删 body,否则会留下指向已删关节的悬空执行器
    for act in list(spec.actuators):
      if act.name.startswith(f"hand_{side}_"):
        spec.delete(act)
    for eq in list(spec.equalities):
      if (eq.name1 or "").startswith(f"hand_{side}_") or (eq.name2 or "").startswith(f"hand_{side}_"):
        spec.delete(eq)
    spec.delete(hand)

    frame = parent.add_frame(pos=pos, quat=quat)
    frame.attach_body(grip_src.body("gripper"), f"grip_{side}_", "")
    a = spec.add_actuator()
    a.name = f"grip_{side}_act"
    a.target = f"grip_{side}_gripper"
    a.trntype = mujoco.mjtTrn.mjTRN_JOINT
    a.gaintype = mujoco.mjtGain.mjGAIN_FIXED
    a.biastype = mujoco.mjtBias.mjBIAS_AFFINE
    a.gainprm = [SO101_KP] + [0.0] * 9
    a.biasprm = [0.0, -SO101_KP, -SO101_KV] + [0.0] * 7
    a.forcerange = [-SO101_FRC, SO101_FRC]
    a.forcelimited = mujoco.mjtLimited.mjLIMITED_TRUE
    info[side] = {"prefix": f"grip_{side}_"}
  spec.assets = {**(spec.assets or {}), **_so101_assets()}
  return info


def _so101_assets():
  d = SO101_XML.parent / "assets"
  return {f.name: f.read_bytes() for f in d.glob("*")} if d.is_dir() else {}


# ---------------------------------------------------------------- 双 SO-101

SO101_JOINTS = ("shoulder_pan", "shoulder_lift", "elbow_flex",
                "wrist_flex", "wrist_roll")
PAIR = {          # 桌面式布局:两台并排装在台面上,都朝 +x
  "table_top": 0.40, "table_x": 0.15, "table_half": (0.30, 0.40, 0.02),
  "mount_y": 0.22, "obj_x": 0.27,
  "obj_ys": (-0.20, -0.10, 0.0, 0.10, 0.20),
}


def build_so101_pair(cfg: Cfg):
  """从零搭一个场景:地面 + 操作台 + 台面上并排两台 SO-101。

  选这个而不是给 P7 换爪:SO-101 只有 5 个自由度,任务约束正好是
  位置 3 + 爪口朝向 2,IK 良定;而且这套夹爪 + 这个尺寸的物体在 Go2 抓放
  demo 里已经端到端验过(随机位置 12 次成功 10 次),数值可以直接搬。
  """
  spec = mujoco.MjSpec()
  spec.option.timestep = 0.002
  spec.option.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
  spec.option.iterations = 20
  spec.option.ls_iterations = 20
  wb = spec.worldbody

  tex = spec.add_texture(
    name="grid", type=mujoco.mjtTexture.mjTEXTURE_2D,
    builtin=mujoco.mjtBuiltin.mjBUILTIN_CHECKER,
    rgb1=[0.17, 0.20, 0.25], rgb2=[0.12, 0.15, 0.19], width=300, height=300)
  del tex
  mat = spec.add_material(name="grid", texrepeat=[8, 8], reflectance=0.05)
  mat.textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] = "grid"
  wb.add_geom(name="floor", type=mujoco.mjtGeom.mjGEOM_PLANE, size=[0, 0, 0.05],
              material="grid", contype=1, conaffinity=1, condim=3)
  wb.add_light(pos=[0.3, 0, 2.5], dir=[0, 0, -1],
               type=mujoco.mjtLightType.mjLIGHT_DIRECTIONAL)

  P = PAIR
  half, top = P["table_half"], P["table_top"]
  t = wb.add_body(name="table", pos=[P["table_x"], 0, 0])
  t.add_geom(name="table_top", type=mujoco.mjtGeom.mjGEOM_BOX, size=list(half),
             pos=[0, 0, top - half[2]], rgba=[0.55, 0.42, 0.30, 1],
             contype=1, conaffinity=1, condim=3, friction=[1.0, 0.005, 0.0001])
  t.add_geom(name="table_top_vis", type=mujoco.mjtGeom.mjGEOM_BOX,
             size=[half[0] * 1.001, half[1] * 1.001, half[2] * 0.999],
             pos=[0, 0, top - half[2]], rgba=[0.55, 0.42, 0.30, 1],
             contype=0, conaffinity=0, mass=0)
  for sx in (-1, 1):
    for sy in (-1, 1):
      t.add_geom(name=f"table_leg_{sx}_{sy}", type=mujoco.mjtGeom.mjGEOM_BOX,
                 size=[0.02, 0.02, (top - 2 * half[2]) / 2],
                 pos=[sx * (half[0] - 0.03), sy * (half[1] - 0.03),
                      (top - 2 * half[2]) / 2],
                 rgba=[0.38, 0.29, 0.21, 1], contype=0, conaffinity=0)

  for side, sy in (("left", 1), ("right", -1)):
    arm = mujoco.MjSpec.from_file(str(SO101_XML))
    frame = wb.add_frame(pos=[0.0, sy * P["mount_y"], top])
    spec.attach(arm, prefix=f"arm_{side}_", frame=frame)
  return spec, {"table": {"x": P["table_x"], "top": top, "half": list(half)}}


def ik_pos(m, d, qadr, rng, eepos, target, iters=400, axis_fn=None, axis_w=None):
  """位置(+可选末端朝向)IK,有限差分雅可比。用来算就绪位形。

  就绪位形必须和前端默认的朝向约束**用同一套约束**解,否则页面一打开
  就在跟自己较劲(实测残差 69mm,手臂一直在往回摆)。
  """
  W = 0.25
  q = np.zeros(len(qadr))
  rows = 6 if axis_fn is not None else 3

  def err(dd):
    e = [target - eepos(dd)]
    if axis_fn is not None:
      e.append(W * (np.asarray(axis_w) - axis_fn(dd)))
    return np.concatenate(e)

  for _ in range(iters):
    d.qpos[qadr] = np.clip(q, rng[:, 0], rng[:, 1])
    mujoco.mj_kinematics(m, d)
    e = err(d)
    if np.linalg.norm(e[:3]) < 5e-4 and (rows == 3 or np.linalg.norm(e[3:]) / W < 0.05):
      break
    J = np.zeros((rows, len(qadr)))
    for j in range(len(qadr)):
      qh = q.copy()
      qh[j] = np.clip(qh[j] + 1e-4, rng[j, 0], rng[j, 1])
      dq = qh[j] - q[j]
      if abs(dq) < 1e-12:
        continue
      d.qpos[qadr] = np.clip(qh, rng[:, 0], rng[:, 1])
      mujoco.mj_kinematics(m, d)
      J[:, j] = (err(d) - e) / -dq
    q = np.clip(q + 0.5 * (J.T @ np.linalg.solve(J @ J.T + 1e-4 * np.eye(rows), e)),
                rng[:, 0], rng[:, 1])
  d.qpos[qadr] = np.clip(q, rng[:, 0], rng[:, 1])
  mujoco.mj_kinematics(m, d)
  return q, float(np.linalg.norm(target - eepos(d)))


# ---------------------------------------------------------------- 就绪位形

ARM_PREFIX = {"left": "L", "right": "R"}   # 左臂关节叫 L1..L7,右臂叫 R1..R7


def arm_joints(side):
  return [f"arm_{side}_{ARM_PREFIX[side]}{i}_Joint" for i in range(1, 8)]


def solve_ready(m, d, side, target):
  """把某一侧的腕部 site 解到 target,返回 7 个臂关节角。"""
  jn = arm_joints(side)
  qadr = np.array([m.jnt_qposadr[m.joint(n).id] for n in jn])
  dadr = np.array([m.jnt_dofadr[m.joint(n).id] for n in jn])
  rng = np.array([m.jnt_range[m.joint(n).id] for n in jn])
  sid = m.site(f"arm_{side}_tool0").id
  q = np.zeros(7)
  jacp, jacr = np.zeros((3, m.nv)), np.zeros((3, m.nv))
  for _ in range(300):
    d.qpos[qadr] = q
    mujoco.mj_kinematics(m, d)
    mujoco.mj_comPos(m, d)
    e = target - d.site_xpos[sid]
    if np.linalg.norm(e) < 5e-4:
      break
    mujoco.mj_jacSite(m, d, jacp, jacr, sid)
    J = jacp[:, dadr]
    q = np.clip(q + 0.6 * (J.T @ np.linalg.solve(J @ J.T + 1e-4 * np.eye(3), e)),
                rng[:, 0], rng[:, 1])
  return q.tolist(), qadr.tolist(), float(np.linalg.norm(target - d.site_xpos[sid]))


def build_p7(cfg: Cfg):
  """linker 双臂工作站 + 操作台 + 物体。"""
  src = WS_ROOT / cfg.workstation / "workstation.mjcf"
  assert src.exists(), f"找不到工作站:{src}"
  # MjSpec 不认 .mjcf 后缀,也不会跟着 ../../ 去找网格。
  # 读成字符串 + 把引用到的网格按**原样的相对路径**塞进 assets 字典。
  import re as _re
  text = src.read_text()
  assets = {}
  for ref in sorted(set(_re.findall(r'file="([^"]+)"', text))):
    f = (src.parent / ref).resolve()
    if f.is_file():
      assets[ref] = f.read_bytes()
    else:
      print(f"  [警告] 网格找不到:{ref}")
  print(f"  载入网格 {len(assets)} 个")
  spec = mujoco.MjSpec.from_string(text, assets=assets)
  if cfg.gripper == "so101":
    swap_gripper(spec)
    print("  末端已换成 SO-101 二指夹爪")
  info = add_scene(spec, cfg)
  spec.option.timestep = 0.002
  spec.option.integrator = mujoco.mjtIntegrator.mjINT_IMPLICITFAST
  spec.option.iterations = 20
  spec.option.ls_iterations = 20
  return spec, info


def main() -> None:
  cfg = tyro.cli(Cfg)
  out = HERE / "scenes" / cfg.name
  out.mkdir(parents=True, exist_ok=True)

  if cfg.rig == "so101_pair":
    spec, info = build_so101_pair(cfg)
    info["objects"] = add_objects(spec, spec.worldbody, PAIR["table_top"],
                                  PAIR["obj_x"], list(PAIR["obj_ys"]))
    print(f"  双 SO-101:台面 z={PAIR['table_top']} 底座 y=±{PAIR['mount_y']}")
  else:
    spec, info = build_p7(cfg)
  m = spec.compile()
  d = mujoco.MjData(m)

  # 就绪位形:末端停在台面上方
  ready = {}
  if cfg.rig == "so101_pair":
    P = PAIR
    for side, sy in (("left", 1), ("right", -1)):
      jn = [f"arm_{side}_{j}" for j in SO101_JOINTS]
      qadr = np.array([m.jnt_qposadr[m.joint(n).id] for n in jn])
      rng = np.array([m.jnt_range[m.joint(n).id] for n in jn])
      ga = [g for g in range(m.ngeom)
            if (m.geom(g).name or "").startswith(f"arm_{side}_fixed_jaw_sph_tip")]
      gb = [g for g in range(m.ngeom)
            if (m.geom(g).name or "").startswith(f"arm_{side}_moving_jaw_sph_tip")]

      def jaw_mid(dd, ga=ga, gb=gb):
        return (dd.geom_xpos[ga].mean(0) + dd.geom_xpos[gb].mean(0)) / 2

      sid = m.site(f"arm_{side}_gripperframe").id

      def jaw_axis(dd, sid=sid):
        return dd.site_xmat[sid].reshape(3, 3)[:, 0]     # 爪口进近轴 = site 局部 +x

      # 爪口朝下的可达高度有限:离台面 6cm 残差 0.3mm、10cm 9mm、14cm 就 38mm 了
      tgt = np.array([0.22, sy * (P["mount_y"] - 0.05), P["table_top"] + 0.07])
      q, err = ik_pos(m, d, qadr, rng, jaw_mid, tgt,
                      axis_fn=jaw_axis, axis_w=[0, 0, -1])
      ready[side] = {"q": q.tolist(), "qadr": qadr.tolist(), "err": err}
      print(f"  就绪位形 {side:5s} 残差 {err * 1000:.1f} mm")
  else:
    for side, ty in (("left", 0.18), ("right", -0.18)):
      q, qadr, err = solve_ready(m, d, side,
                                 np.array([cfg.table_x - 0.06, ty, cfg.table_top + 0.16]))
      ready[side] = {"q": q, "qadr": qadr, "err": err}
      print(f"  就绪位形 {side:5s} 残差 {err * 1000:.1f} mm")

  qpos0 = np.array(m.qpos0, dtype=float).copy()
  for side in ("left", "right"):
    qpos0[ready[side]["qadr"]] = ready[side]["q"]
  key = spec.add_key()
  key.name = "ready"
  key.qpos = qpos0.tolist()
  key.ctrl = [0.0] * m.nu

  inline_meshes(spec, m, cfg.decimate)
  xml = spec.to_xml().replace(' class=""', "")
  (out / "model.xml").write_text(xml)
  mb = (out / "model.xml").stat().st_size / 1e6

  # 重新编译一遍,确认自包含且几何没跑偏,同时取出前端要的索引
  m2 = mujoco.MjModel.from_xml_path(str(out / "model.xml"))
  d2 = mujoco.MjData(m2)
  mujoco.mj_resetDataKeyframe(m2, d2, 0)
  mujoco.mj_forward(m2, d2)

  def act_idx(prefix):
    return [i for i in range(m2.nu) if m2.actuator(i).name.startswith(prefix)]

  arms, hands = {}, {}
  gname = lambda g: m2.geom(g).name or ""            # noqa: E731
  for side in ("left", "right"):
    if cfg.rig == "so101_pair":
      pre = f"arm_{side}_"
      a = [i for i in range(m2.nu)
           if m2.actuator(i).name in [f"{pre}{j}" for j in SO101_JOINTS]]
      arms[side] = {
        "act": a,
        "names": [m2.actuator(i).name for i in a],
        "qadr": [int(m2.jnt_qposadr[int(m2.actuator_trnid[i][0])]) for i in a],
        "range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in a],
        "site": f"{pre}gripperframe",
        "site_id": int(m2.site(f"{pre}gripperframe").id),
        "ready": ready[side]["q"],
      }
      gi = [i for i in range(m2.nu) if m2.actuator(i).name == f"{pre}gripper"]
      hands[side] = {
        "kind": "so101",
        "act": gi,
        "names": [m2.actuator(i).name for i in gi],
        "range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in gi],
        "geoms": [g for g in range(m2.ngeom)
                  if (m2.body(int(m2.geom_bodyid[g])).name or "").startswith(pre)
                  and (m2.geom_contype[g] or m2.geom_conaffinity[g])],
        # 夹持中心 = 两爪指尖球的中点(Go2 抓放 demo 验证过的算法)
        "jaw_a": [g for g in range(m2.ngeom) if gname(g).startswith(f"{pre}fixed_jaw_sph_tip")],
        "jaw_b": [g for g in range(m2.ngeom) if gname(g).startswith(f"{pre}moving_jaw_sph_tip")],
        "palm_body": int(m2.body(f"{pre}gripper").id),
        "site_id": int(m2.site(f"{pre}gripperframe").id),
      }
      continue

    a = act_idx(f"arm_{side}_")
    h = act_idx(f"hand_{side}_")
    arms[side] = {
      "act": a,
      "names": [m2.actuator(i).name for i in a],
      "qadr": [int(m2.jnt_qposadr[int(m2.actuator_trnid[i][0])]) for i in a],
      "range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in a],
      "site": f"arm_{side}_tool0",
      "site_id": int(m2.site(f"arm_{side}_tool0").id),
      "ready": ready[side]["q"],
    }
    if cfg.gripper == "so101":
      pre = f"grip_{side}_"
      h = [i for i in range(m2.nu) if m2.actuator(i).name.startswith(pre)]
      gname = lambda g: m2.geom(g).name or ""      # noqa: E731
      hands[side] = {
        "kind": "so101",
        "act": h,
        "names": [m2.actuator(i).name for i in h],
        "range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in h],
        "geoms": [g for g in range(m2.ngeom)
                  if (m2.body(int(m2.geom_bodyid[g])).name or "").startswith(pre)
                  and (m2.geom_contype[g] or m2.geom_conaffinity[g])],
        # 夹持中心 = 两爪指尖球的中点(Go2 上验证过的算法)
        "jaw_a": [g for g in range(m2.ngeom) if gname(g).startswith(f"{pre}fixed_jaw_sph_tip")],
        "jaw_b": [g for g in range(m2.ngeom) if gname(g).startswith(f"{pre}moving_jaw_sph_tip")],
        "palm_body": int(m2.body(f"{pre}gripper").id),
        "site": f"{pre}gripperframe",
        "site_id": int(m2.site(f"{pre}gripperframe").id),
      }
    else:
      pre = f"hand_{side}_{'l' if side == 'left' else 'r'}h_"
      hands[side] = {
        "kind": "l6",
        "act": h,
        "names": [m2.actuator(i).name for i in h],
        "range": [[float(x) for x in m2.actuator_ctrlrange[i]] for i in h],
        # 原始资产的 geom 是靠 class 配的、大多没名字,只能按所属 body 来筛
        "geoms": [g for g in range(m2.ngeom)
                  if (m2.body(int(m2.geom_bodyid[g])).name or "").startswith(f"hand_{side}_")
                  and (m2.geom_contype[g] or m2.geom_conaffinity[g])],
        "tip_bodies": [int(m2.body(f"{pre}{t}_tip").id)
                       for t in ("thumb", "index", "middle", "ring", "pinky")],
        "palm_body": int(m2.body(f"{pre}hand_base_link").id),
      }

  meta = {
    "name": cfg.name,
    "workstation": cfg.workstation,
    "gripper": cfg.gripper,
    "sim": {"timestep": float(m2.opt.timestep), "substeps": 10},
    "nq": int(m2.nq), "nu": int(m2.nu), "ngeom": int(m2.ngeom),
    "kps": [float(m2.actuator_gainprm[i][0]) for i in range(m2.nu)],
    "arms": arms, "hands": hands,
    "scene": info,
    "xmlMB": round(mb, 2),
  }
  # 物体的 qpos 地址(自由关节),前端要用来改位置
  for o in meta["scene"]["objects"]:
    bid = m2.body(o["body"]).id
    o["qadr"] = int(m2.jnt_qposadr[m2.body_jntadr[bid]])
    o["body_id"] = int(bid)
    o["geom_id"] = int(m2.geom(o["geom"]).id)
  (out / "meta.json").write_text(json.dumps(meta, ensure_ascii=False))

  print(f"已导出 {cfg.name} → {out}")
  print(f"  model.xml {mb:.2f} MB | nq={m2.nq} nu={m2.nu} ngeom={m2.ngeom} nmesh={m2.nmesh}")
  print(f"  桌面 z={info['table']['top']} 中心 x={info['table']['x']} | 物体 "
        + ", ".join(o["name"] for o in meta["scene"]["objects"]))


if __name__ == "__main__":
  main()