const decoder=new TextDecoder('utf-8');
const encoder=new TextEncoder();

function numbers(value:string|undefined,count:number):number[]{
  const parsed=(value??'').trim().split(/\s+/).filter(Boolean).map(Number);
  return Array.from({length:count},(_,index)=>Number.isFinite(parsed[index])?parsed[index]:0);
}

export type UrdfBaseMode='floating'|'fixed';
export type CameraDirection='+X'|'-X'|'+Y'|'-Y'|'+Z'|'-Z';
export interface UrdfEnhancementOptions {addActuators:boolean;addSensors:boolean;sensorType:'camera';cameraMountBody?:string;cameraPosition?:[number,number,number];cameraDirection?:CameraDirection;}
export interface UrdfEnhancementResult {data:Uint8Array;actuatorCount:number;cameraAdded:boolean;imuAdded:boolean;unitreeGo2wTuned:boolean;}

function uniqueName(document:Document,selector:string,base:string):string {
  const occupied=new Set(Array.from(document.querySelectorAll(selector)).map(element=>element.getAttribute('name')).filter(Boolean));
  if(!occupied.has(base))return base;
  for(let index=2;;index+=1)if(!occupied.has(`${base}_${index}`))return `${base}_${index}`;
}

function tuneUnitreeGo2w(document:Document):boolean {
  const prefixes=['FL','FR','RL','RR'],legParts=['hip','thigh','calf'];
  const expected=[...prefixes.flatMap(prefix=>legParts.map(part=>`${prefix}_${part}_joint`)),...prefixes.map(prefix=>`${prefix}_foot_joint`)];
  const joints=new Map(Array.from(document.querySelectorAll('worldbody joint[name]')).map(joint=>[joint.getAttribute('name')!,joint]));
  if(!expected.every(name=>joints.has(name)))return false;
  const root=document.querySelector('mujoco');if(!root)return false;
  let option=Array.from(root.children).find(element=>element.tagName==='option');if(!option){option=document.createElement('option');root.prepend(option);}
  option.setAttribute('cone','elliptic');option.setAttribute('impratio','100');
  for(const name of expected){const joint=joints.get(name)!;joint.setAttribute('armature','0.01');joint.setAttribute('damping','0.1');joint.setAttribute('frictionloss','0.2');}
  for(const motor of Array.from(document.querySelectorAll('actuator > *[joint]'))){const jointName=motor.getAttribute('joint')??'';if(!expected.includes(jointName))continue;const limit=jointName.includes('_calf_')?45.43:jointName.endsWith('_foot_joint')?15:23.7;motor.setAttribute('forcelimited','true');motor.setAttribute('forcerange',`${-limit} ${limit}`);}
  for(const geom of Array.from(document.querySelectorAll('worldbody body geom'))){if(geom.getAttribute('contype')==='0'||geom.getAttribute('group')==='1')continue;geom.setAttribute('friction','0.4 0.005 0.0001');geom.setAttribute('margin','0.001');geom.setAttribute('condim','1');const body=geom.closest('body'),name=body?.getAttribute('name')??'';if(/_(?:foot|wheel)(?:_link)?$/i.test(name)){geom.setAttribute('friction','0.8 0.02 0.01');geom.setAttribute('condim','6');geom.setAttribute('priority','1');}}
  return true;
}

/** 为转换后的 MJCF 按需补充可直接控制的关节电机和一台自动取景相机。 */
export function enhanceConvertedMjcf(data:Uint8Array,options:UrdfEnhancementOptions):UrdfEnhancementResult {
  const document=new DOMParser().parseFromString(decoder.decode(data),'application/xml');
  if(document.querySelector('parsererror'))throw new Error('MuJoCo 导出的 MJCF XML 无法解析');
  const root=document.querySelector('mujoco'),worldbody=document.querySelector('mujoco > worldbody');
  if(!root||!worldbody)throw new Error('MuJoCo 导出的 MJCF 缺少 worldbody');
  let actuatorCount=0;
  if(options.addActuators){
    let actuator=Array.from(root.children).find(element=>element.tagName==='actuator');
    if(!actuator){actuator=document.createElement('actuator');root.append(actuator);}
    const driven=new Set(Array.from(actuator.children).flatMap(element=>[element.getAttribute('joint'),element.getAttribute('jointinparent')]).filter(Boolean));
    for(const joint of Array.from(worldbody.querySelectorAll('joint'))){
      const type=joint.getAttribute('type')??'hinge';
      if(type!=='hinge'&&type!=='slide')continue;
      let jointName=joint.getAttribute('name');
      if(!jointName){jointName=uniqueName(document,'joint[name]','__platform_joint__');joint.setAttribute('name',jointName);}
      if(driven.has(jointName))continue;
      const motor=document.createElement('motor');
      motor.setAttribute('name',uniqueName(document,'actuator > *[name]',`${jointName}_motor`));
      motor.setAttribute('joint',jointName);motor.setAttribute('ctrllimited','false');motor.setAttribute('forcelimited','true');motor.setAttribute('forcerange','-100 100');
      if(!joint.hasAttribute('stiffness'))joint.setAttribute('stiffness','0');if(!joint.hasAttribute('damping'))joint.setAttribute('damping','0');
      actuator.append(motor);driven.add(jointName);actuatorCount+=1;
    }
    if(!actuator.children.length)actuator.remove();
  }
  const unitreeGo2wTuned=tuneUnitreeGo2w(document);
  let imuAdded=false;
  if(options.addSensors&&!document.querySelector('sensor > gyro[name="imu_gyro"], sensor > gyro[name="__platform_imu_gyro__"]')){
    const bodies=Array.from(worldbody.querySelectorAll('body')),mount=bodies.find(body=>/^(?:base|base_link|trunk)$/i.test(body.getAttribute('name')??''))??Array.from(worldbody.children).find(element=>element.tagName==='body');
    if(mount){
      const site=document.createElement('site'),siteName=uniqueName(document,'site[name]','imu');site.setAttribute('name',siteName);site.setAttribute('pos','0 0 0');site.setAttribute('size','0.005');site.setAttribute('rgba','0.2 0.8 1 0.5');mount.append(site);
      let sensor=Array.from(root.children).find(element=>element.tagName==='sensor');if(!sensor){sensor=document.createElement('sensor');root.append(sensor);}
      const gyro=document.createElement('gyro');gyro.setAttribute('name',uniqueName(document,'sensor > *[name]','imu_gyro'));gyro.setAttribute('site',siteName);gyro.setAttribute('noise','0');sensor.append(gyro);
      const accelerometer=document.createElement('accelerometer');accelerometer.setAttribute('name',uniqueName(document,'sensor > *[name]','imu_acc'));accelerometer.setAttribute('site',siteName);accelerometer.setAttribute('noise','0');sensor.append(accelerometer);imuAdded=true;
    }
  }
  let cameraAdded=false;
  if(options.addSensors&&options.sensorType==='camera'&&!document.querySelector('camera[name^="__platform_camera__"]')){
    const bodies=Array.from(worldbody.querySelectorAll('body')),preferred=options.cameraMountBody?bodies.find(body=>body.getAttribute('name')===options.cameraMountBody):undefined;
    const mount=preferred??bodies.find(body=>/(head|camera|sensor|neck|头)/i.test(body.getAttribute('name')??''))??bodies.at(-1);
    if(mount){
      const directions:Record<CameraDirection,string>={'+X':'0 -1 0 0 0 1','-X':'0 1 0 0 0 1','+Y':'1 0 0 0 0 1','-Y':'-1 0 0 0 0 1','+Z':'0 1 0 1 0 0','-Z':'0 -1 0 1 0 0'},position=options.cameraPosition??[.1,0,.05],direction=options.cameraDirection??'+X';
      const camera=document.createElement('camera');camera.setAttribute('name',uniqueName(document,'camera[name]','__platform_camera__'));camera.setAttribute('mode','fixed');
      camera.setAttribute('pos',position.map(value=>Number.isFinite(value)?value:0).join(' '));camera.setAttribute('xyaxes',directions[direction]);camera.setAttribute('fovy','45');camera.setAttribute('resolution','640 480');mount.append(camera);cameraAdded=true;
    }
  }
  return {data:encoder.encode(new XMLSerializer().serializeToString(document)),actuatorCount,cameraAdded,imuAdded,unitreeGo2wTuned};
}

/** 给 MuJoCo 从 URDF 导出的 MJCF 添加地面、设置基座类型，并整体抬升根 body。 */
export function groundConvertedMjcf(data:Uint8Array,minimumZ:number,baseMode:UrdfBaseMode='fixed'):Uint8Array {
  const document=new DOMParser().parseFromString(decoder.decode(data),'application/xml');
  if(document.querySelector('parsererror'))throw new Error('MuJoCo 导出的 MJCF XML 无法解析');
  const worldbody=document.querySelector('mujoco > worldbody');
  if(!worldbody)throw new Error('MuJoCo 导出的 MJCF 缺少 worldbody');
  const lift=Number.isFinite(minimumZ)?-minimumZ:0;
  const rootBodies=Array.from(worldbody.children).filter(element=>element.tagName==='body');
  for(const [index,body] of rootBodies.entries()){
    const pos=numbers(body.getAttribute('pos')??undefined,3);pos[2]+=lift;body.setAttribute('pos',pos.join(' '));
    if(baseMode==='floating'&&!Array.from(body.children).some(element=>element.tagName==='freejoint'||element.tagName==='joint')){
      const freejoint=document.createElement('freejoint');freejoint.setAttribute('name',rootBodies.length===1?'__platform_base_freejoint__':`__platform_base_freejoint_${index}__`);body.prepend(freejoint);
    }
  }
  const floor=document.createElement('geom');
  floor.setAttribute('name','__platform_ground__');floor.setAttribute('type','plane');floor.setAttribute('size','1 1 0.1');floor.setAttribute('pos','0 0 0');floor.setAttribute('rgba','0.12 0.16 0.22 1');floor.setAttribute('group','5');floor.setAttribute('friction','1 0.005 0.0001');
  worldbody.prepend(floor);
  return encoder.encode(new XMLSerializer().serializeToString(document));
}