import loadMujoco, {type MainModule} from '@mujoco/mujoco';
import type {ProjectManifest} from '../project/types';
import {prepareProjectForMujoco} from '../project/importer';
import {groundConvertedMjcf,type UrdfBaseMode} from '../project/urdfToMjcf';
import {MemfsWorkspace} from '../project/workspace';
import {SimulationSession, type FrameResult, type SimulationSnapshot} from './SimulationSession';

export type UrdfLoadMode='mjcf'|'native';
export type {UrdfBaseMode};

export interface PhysicsAdapter {
  load(manifest:ProjectManifest,entryPath:string,urdfMode?:UrdfLoadMode,baseMode?:UrdfBaseMode):Promise<SimulationSnapshot>;
  advance(now: number): FrameResult;
  snapshot(): SimulationSnapshot | null;
  setPaused(paused: boolean): void;
  setSpeed(speed: number): void;
  reset(): void;
  singleStep(): void;
  setActuator(id: number, value: number): void;
  setJointPosition(id:number,value:number):boolean;
  resetJoints():void;
  setIgnoreJointLimits(ignore:boolean):void;
  setExternalForce(bodyId: number, force: [number, number, number]): void;
  clearExternalForce(): void;
  dispose(): void;
}

let modulePromise: Promise<MainModule> | undefined;
export function getMujocoModule(): Promise<MainModule> {
  if (!modulePromise) {
    console.info('[MuJoCo] 开始初始化单线程 WASM');
    modulePromise = loadMujoco().then((module) => {
      console.info('[MuJoCo] WASM 初始化完成');
      return module;
    });
  }
  return modulePromise;
}

export class MainThreadPhysicsAdapter implements PhysicsAdapter {
  session: SimulationSession | null = null;
  workspace: MemfsWorkspace | null = null;

  async load(manifest:ProjectManifest,entryPath:string,urdfMode:UrdfLoadMode='mjcf',baseMode:UrdfBaseMode='floating'):Promise<SimulationSnapshot> {
    this.releaseCurrent();const module=await getMujocoModule();const workspace=new MemfsWorkspace(module,manifest.id);const prepared=prepareProjectForMujoco(manifest,entryPath);
    try{
      console.info('[MuJoCo] 写入 MEMFS',prepared.manifest.files.length);workspace.mount(prepared.manifest);
      const entry=prepared.manifest.entries.find(candidate=>candidate.path===entryPath);let modelPath=workspace.path(entryPath);const warnings=[...prepared.warnings];
      if(entry?.format==='urdf'&&urdfMode==='mjcf'){
        console.info('[MuJoCo] 编译 URDF 中间模型',entryPath);const intermediate=new SimulationSession(module,modelPath);
        try{
          const minimumZ=intermediate.minimumGeometryZ();const slash=entryPath.lastIndexOf('/');const directory=slash>=0?entryPath.slice(0,slash+1):'';const convertedPath=`${directory}.__mujoco_converted_${manifest.id.replace(/[^a-zA-Z0-9_-]/g,'_')}.xml`;
          if(module.mj_saveLastXML(workspace.path(convertedPath),intermediate.model)===0)throw new Error('MuJoCo 无法导出中间 MJCF');
          workspace.writeGenerated(convertedPath,groundConvertedMjcf(new TextEncoder().encode(workspace.readText(convertedPath)),minimumZ,baseMode));modelPath=workspace.path(convertedPath);
          warnings.push(`URDF 已转换为 MJCF（${baseMode==='floating'?'浮动基座':'固定基座'}），并整体平移 ${(-minimumZ).toFixed(4)} m，使最低点接触 z=0 地面`);
        }finally{intermediate.dispose();}
      }
      console.info('[MuJoCo] 编译模型',modelPath);const session=new SimulationSession(module,modelPath,warnings);
      if(entry?.format==='urdf'&&urdfMode==='native'){const offset=session.alignLowestPointToGround();warnings.push(`原生 URDF 已整体平移 ${offset.toFixed(4)} m，使最低点位于 z=0`);}
      if(warnings.length)console.warn('[MuJoCo] URDF 兼容处理',warnings);
      console.info('[MuJoCo] 模型编译完成');this.workspace=workspace;this.session=session;const snapshot=session.snapshot();console.info('[MuJoCo] 状态快照完成');return snapshot;
    }catch(error){workspace.dispose();throw new Error(`模型编译失败（${entryPath}）：${error instanceof Error?error.message:String(error)}`,{cause:error});}
  }
  advance(now:number):FrameResult{return this.session?.advance(now)??{steps:0,stepMs:0,overBudget:false};}
  snapshot():SimulationSnapshot|null{return this.session?.snapshot()??null;}
  setPaused(value:boolean):void{this.session?.setPaused(value);}
  setSpeed(value:number):void{this.session?.setSpeed(value);}
  reset():void{this.session?.reset();}
  singleStep():void{this.session?.singleStep();}
  setActuator(id:number,value:number):void{this.session?.setActuator(id,value);}
  setJointPosition(id:number,value:number):boolean{return this.session?.setJointPosition(id,value)??false;}
  resetJoints():void {this.session?.resetJoints();}
  setIgnoreJointLimits(ignore:boolean):void {this.session?.setIgnoreJointLimits(ignore);}
  setExternalForce(bodyId:number,force:[number,number,number]):void{this.session?.setExternalForce(bodyId,force);}
  clearExternalForce():void{this.session?.clearExternalForce();}
  private releaseCurrent():void{this.session?.dispose(); this.session=null; this.workspace?.dispose(); this.workspace=null;}
  dispose():void{this.releaseCurrent();}
}