import type { ProjectManifest } from '../project/types';
import { MOBILE_TASK as T, type RobotConfig } from './RobotDescriptor';

const encoder = new TextEncoder();
function parse(data: Uint8Array): Document {
  const doc = new DOMParser().parseFromString(new TextDecoder().decode(data), 'application/xml');
  if (doc.querySelector('parsererror') || doc.doctype) throw new Error('模型 XML 无效');
  return doc;
}
function node(
  doc: Document,
  tag: string,
  attributes: Record<string, string | number> = {},
): Element {
  const n = doc.createElement(tag);
  for (const [key, value] of Object.entries(attributes)) n.setAttribute(key, String(value));
  return n;
}
function named(doc: Document, tag: string, name: string): Element {
  const result = Array.from(doc.getElementsByTagName(tag)).filter(
    (e) => e.getAttribute('name') === name,
  );
  if (result.length !== 1) throw new Error(`模型缺少/重复 ${tag}: ${name}`);
  return result[0];
}

/** The new bundle has a different arm. Never route it through the old source/FK gate. */
export function prepareBundle(
  manifest: ProjectManifest,
  entryPath: string,
  c: RobotConfig,
): ProjectManifest {
  const source = manifest.files.find((f) => f.path === entryPath);
  if (!source) throw new Error('缺少 bundle URDF');
  const doc = parse(source.data);
  for (const name of [...c.baseJoints, ...c.armJoints.map((j) => j.name), c.gripperJoint])
    named(doc, 'joint', name);
  for (const mesh of doc.querySelectorAll('mesh[filename]')) {
    if (mesh.getAttribute('filename')!.includes('4-Omni-Directional-Wheel_Single_Body'))
      mesh.replaceWith(node(doc, 'sphere', { radius: 0.05 }));
  }
  // Exporter duplicated collisions. Keep one representation per visual, not overlapping copies.
  for (const link of doc.querySelectorAll('link')) {
    const seen = new Set<string>();
    for (const collision of link.querySelectorAll(':scope > collision')) {
      const mesh = collision.querySelector('mesh');
      if (!mesh) continue;
      const origin = collision.querySelector('origin');
      const transform = [
        mesh.getAttribute('scale') ?? '1 1 1',
        origin?.getAttribute('xyz') ?? '0 0 0',
        origin?.getAttribute('rpy') ?? '0 0 0',
      ]
        .map((text) =>
          text
            .trim()
            .split(/\s+/)
            .map((v) => Math.round(Number(v) * 1e9) / 1e9)
            .join(' '),
        )
        .join('|');
      const key = `${mesh.getAttribute('filename')}|${transform}`;
      if (seen.has(key)) collision.remove();
      else seen.add(key);
    }
  }
  const data = encoder.encode(new XMLSerializer().serializeToString(doc));
  return {
    ...manifest,
    files: manifest.files.map((f) => (f === source ? { ...f, data, size: data.length } : f)),
  };
}

/** Simulation-only ideal omni base; preserves Link1..4 arm frames/inertias/limits.
 * Native arm mesh contacts are convex approximations, not the old CAD hull recipe.
 */
function adaptBundle(doc: Document, c: RobotConfig): void {
  const root = doc.documentElement,
    cad = named(doc, 'body', c.baseBodyName);
  for (const name of c.baseJoints) named(doc, 'joint', name).parentElement!.remove();
  for (const j of cad.querySelectorAll(':scope > joint, :scope > freejoint')) j.remove();
  cad.setAttribute('name', '__mm_bundle_cad');
  cad.setAttribute('pos', '0 0 -0.01786');
  cad.setAttribute('quat', '0.7071067811865476 0 0 -0.7071067811865476');
  const base = node(doc, 'body', { name: c.baseBodyName, pos: '0 0 0.0515' });
  base.append(node(doc, 'freejoint', { name: c.baseJointName }));
  base.append(node(doc, 'inertial', { pos: '0 0 .04', mass: 2.2, diaginertia: '.014 .014 .022' }));
  cad.replaceWith(base);
  base.append(cad);
  // Keep collision on the arm, simplify chassis to avoid convex CAD covering wheel wells.
  const arm = named(doc, 'body', 'base_link');
  for (const geom of cad.querySelectorAll('geom')) {
    if (!arm.contains(geom)) {
      geom.setAttribute('contype', '0');
      geom.setAttribute('conaffinity', '0');
    } else {
      geom.setAttribute('friction', '.8 .005 .0001');
    }
  }
  base.append(
    node(doc, 'geom', { type: 'cylinder', size: '.112 .025', pos: '0 0 .04', group: 3, mass: 0 }),
  );
  root.querySelector('actuator')?.remove();
  const actuators = node(doc, 'actuator');
  root.append(actuators);
  const specs = [
    ...c.armJoints,
    { name: c.gripperJoint, min: c.gripperClosed, max: c.gripperOpen, mode: 'position' },
  ];
  for (const spec of specs) {
    const joint = named(doc, 'joint', spec.name);
    joint.setAttribute('limited', 'true');
    joint.setAttribute('range', `${spec.min} ${spec.max}`);
    joint.setAttribute('damping', '.02');
    joint.setAttribute('armature', '.002');
    // URDF conversion leaves a joint-level ±1 N·m clamp. It otherwise overrides
    // the simulation servos below and lets the loaded arm collapse under gravity.
    joint.setAttribute('actuatorfrclimited', 'true');
    joint.setAttribute('actuatorfrcrange', spec.name === c.gripperJoint ? '-2 2' : '-8 8');
    if (spec.name === c.gripperJoint)
      joint.setAttribute(
        'axis',
        joint
          .getAttribute('axis')!
          .split(/\s+/)
          .map((v) => -Number(v))
          .join(' '),
      );
    actuators.append(
      node(doc, 'position', {
        name: `${spec.name}_servo`,
        joint: spec.name,
        kp: spec.name === c.gripperJoint ? 15 : 40,
        kv: 2,
        ctrllimited: 'true',
        ctrlrange: `${spec.min} ${spec.max}`,
        forcelimited: 'true',
        forcerange: spec.name === c.gripperJoint ? '-2 2' : '-8 8',
      }),
    );
  }
  for (let i = 0; i < c.baseJoints.length; i++) {
    const name = c.baseJoints[i],
      tx = c.baseMix[i][0] * 0.05,
      ty = c.baseMix[i][1] * 0.05;
    const axle = `${-ty} ${tx} 0`;
    const wheel = node(doc, 'body', {
      name: `__mm_${name}`,
      pos: `${0.125 * ty} ${-0.125 * tx} 0`,
    });
    wheel.append(
      node(doc, 'joint', {
        name,
        axis: axle,
        limited: 'false',
        damping: '.0002',
        armature: '.00005',
      }),
    );
    wheel.append(
      node(doc, 'inertial', { pos: '0 0 0', mass: '.06', diaginertia: '.00004 .00004 .00004' }),
    );
    wheel.append(
      node(doc, 'geom', {
        type: 'cylinder',
        size: '.033 .012',
        zaxis: axle,
        contype: 0,
        conaffinity: 0,
        mass: 0,
        group: 1,
        rgba: '.15 .18 .2 1',
      }),
    );
    for (let j = 0; j < 12; j++) {
      const angle = (2 * Math.PI * j) / 12,
        s = Math.sin(angle),
        k = Math.cos(angle);
      const axis = `${k * tx} ${k * ty} ${-s}`;
      const roller = node(doc, 'body', { pos: `${0.041 * s * tx} ${0.041 * s * ty} ${0.041 * k}` });
      roller.append(
        node(doc, 'joint', { axis, limited: 'false', damping: '.000005', armature: '.0000001' }),
      );
      roller.append(
        node(doc, 'geom', {
          type: 'capsule',
          size: '.009 .006',
          zaxis: axis,
          mass: '.003',
          contype: 4,
          conaffinity: 1,
          group: 1,
          friction: '1 .001 .0001',
          condim: 3,
          solref: '.008 1',
          rgba: '.3 .3 .3 1',
        }),
      );
      wheel.append(roller);
    }
    base.append(wheel);
    actuators.append(
      node(doc, 'velocity', {
        name: c.baseActuators[i],
        joint: name,
        kv: '.5',
        ctrllimited: 'true',
        ctrlrange: `${-c.wheelLimit} ${c.wheelLimit}`,
        forcelimited: 'true',
        forcerange: '-2 2',
      }),
    );
  }
}

export function composeMobileScene(data: Uint8Array, config: RobotConfig): Uint8Array {
  const doc = parse(data),
    root = doc.documentElement;
  if (root.tagName !== 'mujoco') throw new Error('场景组合需要已转换的 MJCF');
  if (
    Array.from(doc.querySelectorAll('[name]')).some((e) =>
      e.getAttribute('name')!.startsWith('__mm_'),
    )
  )
    throw new Error('任务命名空间 __mm_ 冲突（请勿重复组合）');
  if (config.recipe === 'lekiwi-bundle') adaptBundle(doc, config);
  const world = root.querySelector('worldbody');
  if (!world) throw new Error('缺少 worldbody');
  // A single task floor, no inherited robot-scene floors.
  for (const geom of world.querySelectorAll(':scope > geom[type="plane"]')) geom.remove();
  world.append(
    node(doc, 'geom', {
      name: '__mm_floor',
      type: 'plane',
      group: 2,
      size: '3 3 .1',
      friction: '.8 .005 .0001',
      contype: 1,
      conaffinity: 7,
      rgba: '.25 .3 .35 1',
    }),
  );
  const object = node(doc, 'body', { name: '__mm_object', pos: T.objectStart.join(' ') });
  object.append(node(doc, 'freejoint', { name: '__mm_object_joint' }));
  object.append(
    node(doc, 'geom', {
      name: '__mm_object_geom',
      type: 'box',
      group: 1,
      size: `${T.objectHalfSize} ${T.objectHalfSize} ${T.objectHalfSize}`,
      mass: '.04',
      friction: '1 .005 .0001',
      condim: 4,
      contype: 1,
      conaffinity: 7,
      rgba: '1 .35 .08 1',
    }),
  );
  world.append(object);
  const goal = node(doc, 'body', { name: '__mm_goal', mocap: 'true', pos: T.goalStart.join(' ') });
  goal.append(
    node(doc, 'geom', {
      name: '__mm_goal_geom',
      type: 'cylinder',
      group: 1,
      size: '.05 .003',
      contype: 0,
      conaffinity: 0,
      rgba: '.1 1 .35 .45',
    }),
  );
  world.append(goal);
  if (config.eefSiteName)
    named(doc, 'body', config.eefBodyName).append(
      node(doc, 'site', {
        name: config.eefSiteName,
        pos: config.eefOffset.join(' '),
        size: '.007',
        rgba: '.2 .8 1 1',
      }),
    );
  const option = root.querySelector('option') ?? node(doc, 'option');
  if (!option.parentElement) root.prepend(option);
  // Preserve reviewed source physics (the full LeKiwi collision recipe uses 1 ms).
  if (!option.hasAttribute('timestep')) option.setAttribute('timestep', '.002');
  if (!option.hasAttribute('integrator')) option.setAttribute('integrator', 'implicitfast');
  root.querySelector('keyframe')?.remove(); // nq grew: old keys are no longer valid.
  return encoder.encode(new XMLSerializer().serializeToString(doc));
}