import type { MainModule } from '@mujoco/mujoco';
import type { ProjectFile, ProjectManifest } from '../project/types';
import { prepareProjectForMujoco } from '../project/importer';
import {
  enhanceConvertedMjcf,
  groundConvertedMjcf,
  type UrdfBaseMode,
  type UrdfEnhancementOptions,
} from '../project/urdfToMjcf';
import { MemfsWorkspace } from '../project/workspace';
import {
  SimulationSession,
  type ActuatorParameters,
  type FrameResult,
  type SimulationSnapshot,
} from './SimulationSession';
import type { ControllerCommand, ControllerStatus } from '../controller/types';
import type { RLCommand, RLPolicyStatus } from '../rl/types';
import type {
  DataRecorderConfig,
  DataRecorderStatus,
  TelemetryChannel,
} from '../telemetry/DataRecorder';
import { composePhysicalMap } from '../map/physicalMap';
import { composeProjectMap } from '../map/MapComposer';
import { resolveProjectMap } from '../map/MapLoader';
import { DEFAULT_MAP_SELECTION, type MapSelection } from '../map/types';

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

export interface PhysicsLoadProgress {
  value: number;
  label: string;
}

export interface PhysicsLoadOptions {
  urdfMode?: UrdfLoadMode;
  baseMode?: UrdfBaseMode;
  enhancements?: UrdfEnhancementOptions;
  map?: MapSelection;
  onProgress?: (progress: PhysicsLoadProgress) => void;
}

export interface PhysicsAdapter {
  load(
    manifest: ProjectManifest,
    entryPath: string,
    options?: PhysicsLoadOptions,
  ): 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;
  setActuatorParameters(id: number, parameters: ActuatorParameters): boolean;
  setJointPosition(id: number, value: number): boolean;
  resetJoints(): void;
  setIgnoreJointLimits(ignore: boolean): void;
  setExternalForce(bodyId: number, force: [number, number, number]): void;
  clearExternalForce(): void;
  loadPythonController(source: string, path: string): Promise<ControllerStatus>;
  setControllerEnabled(enabled: boolean): void;
  sendControllerCommand(command: ControllerCommand): void;
  removeController(): void;
  loadRLPolicy(model: Uint8Array, path: string): Promise<RLPolicyStatus>;
  setRLPolicyEnabled(enabled: boolean): void;
  setRLCommand(command: RLCommand): void;
  removeRLPolicy(): void;
  configureDataRecorder(config: Partial<DataRecorderConfig>): DataRecorderStatus | undefined;
  startDataRecording(): DataRecorderStatus | undefined;
  stopDataRecording(): DataRecorderStatus | undefined;
  clearDataRecording(): DataRecorderStatus | undefined;
  registerDataChannel(channel: TelemetryChannel): (() => void) | undefined;
  exportDataRecording(format: 'csv' | 'json'): Uint8Array;
  cachedSupportFiles(): ProjectFile[];
  releaseRetired(): void;
  rollbackRetired(): void;
  exportMjcf(): Uint8Array;
  dispose(): void;
}

let modulePromise: Promise<MainModule> | undefined;
export function getMujocoModule(): Promise<MainModule> {
  if (!modulePromise) {
    console.info('[MuJoCo] 开始初始化单线程 WASM');
    modulePromise = import('@mujoco/mujoco')
      .then(({ default: loadMujoco }) => loadMujoco())
      .then((module) => {
        console.info('[MuJoCo] WASM 初始化完成');
        return module;
      })
      .catch((error) => {
        // 初始化失败后允许用户重试，而不是永久复用已 rejected 的 Promise。
        modulePromise = undefined;
        throw error;
      });
  }
  return modulePromise;
}

export class MainThreadPhysicsAdapter implements PhysicsAdapter {
  session: SimulationSession | null = null;
  workspace: MemfsWorkspace | null = null;
  private supportFiles: ProjectFile[] = [];
  private loadGeneration = 0;
  private disposed = false;
  private retiredSession: SimulationSession | null = null;
  private retiredWorkspace: MemfsWorkspace | null = null;
  private retiredSupportFiles: ProjectFile[] = [];

  async load(
    manifest: ProjectManifest,
    entryPath: string,
    options: PhysicsLoadOptions = {},
  ): Promise<SimulationSnapshot> {
    if (this.disposed) throw new Error('物理适配器已释放');
    const generation = ++this.loadGeneration;
    const urdfMode = options.urdfMode ?? 'mjcf';
    const baseMode = options.baseMode ?? 'floating';
    const enhancements = options.enhancements ?? {
      addActuators: false,
      addSensors: false,
      sensorType: 'camera',
    };
    const mapSelection = options.map ?? DEFAULT_MAP_SELECTION;
    const report = (value: number, label: string) => {
      if (!this.disposed && generation === this.loadGeneration)
        options.onProgress?.({ value, label });
    };
    report(0.08, '初始化 MuJoCo WebAssembly');
    const module = await getMujocoModule();
    if (this.disposed || generation !== this.loadGeneration) throw new Error('模型加载已取消');
    report(0.32, '校验并准备模型资源');
    const workspace = new MemfsWorkspace(module, `${manifest.id}_stage_${generation}`);
    const prepared = await prepareProjectForMujoco(manifest, entryPath);
    const supportFiles = prepared.manifest.files.filter(
      (file) => !manifest.files.some((original) => original.path === file.path),
    );
    let nextSession: SimulationSession | null = null;
    try {
      report(0.46, '写入浏览器内存文件系统');
      console.info('[MuJoCo] 写入 MEMFS', prepared.manifest.files.length);
      workspace.mount(prepared.manifest);
      const entry = prepared.manifest.entries.find((candidate) => candidate.path === entryPath);
      let modelRelativePath = entryPath,
        modelPath = workspace.path(modelRelativePath);
      const warnings = [...prepared.warnings];
      if (entry?.format === 'urdf' && urdfMode === 'mjcf') {
        report(0.58, '转换并增强 URDF 模型');
        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');
          const grounded = groundConvertedMjcf(
            new TextEncoder().encode(workspace.readText(convertedPath)),
            minimumZ,
            baseMode,
          );
          const enhanced = enhanceConvertedMjcf(grounded, enhancements);
          workspace.writeGenerated(convertedPath, enhanced.data);
          modelRelativePath = convertedPath;
          modelPath = workspace.path(modelRelativePath);
          warnings.push(
            `URDF 已转换为 MJCF（${baseMode === 'floating' ? '浮动基座' : '固定基座'}），并整体平移 ${(-minimumZ).toFixed(4)} m，使最低点接触 z=0 地面`,
          );
          if (enhanced.actuatorCount)
            warnings.push(
              `已为 ${enhanced.actuatorCount} 个 hinge/slide 关节生成 motor 驱动器（控制输入不限幅；hinge 输出单位 N·m，slide 输出单位 N）`,
            );
          if (enhanced.unitreeGo2wTuned)
            warnings.unshift(
              '已识别 Unitree Go2-W，并补齐官方 MuJoCo 关节惯量/阻尼、力矩限幅和轮胎接触参数',
            );
          if (enhanced.imuAdded)
            warnings.unshift(
              '已在浮动基座添加6轴 IMU：imu_gyro（三轴角速度）和 imu_acc（三轴加速度）',
            );
          if (enhanced.cameraAdded)
            warnings.push(
              `已将 640×480 摄像头固连到 ${enhancements.cameraMountBody || '自动选择的头部/末端 body'}，局部位置 ${(enhancements.cameraPosition ?? [0.1, 0, 0.05]).join(' ')} m，朝向 ${enhancements.cameraDirection ?? '+X'}`,
            );
        } finally {
          intermediate.dispose();
        }
      }
      if (mapSelection.kind !== 'none') {
        report(0.72, '组合机器人与物理地图');
        if (entry?.format === 'urdf' && urdfMode === 'native')
          throw new Error('原生 URDF 模式暂不支持地图，请切换为转换模式');
        const slash = modelRelativePath.lastIndexOf('/');
        const directory = slash >= 0 ? modelRelativePath.slice(0, slash + 1) : '';
        const mapPath = `${directory}.__mujoco_map_scene_${manifest.id.replace(/[^a-zA-Z0-9_-]/g, '_')}.xml`;
        if (mapSelection.kind === 'builtin') {
          const composed = composePhysicalMap(
            new TextEncoder().encode(workspace.readText(modelRelativePath)),
            mapSelection.config,
          );
          workspace.writeGenerated(mapPath, composed.data);
          if (composed.summary) warnings.push(composed.summary);
        } else {
          const resolvedMap = resolveProjectMap(prepared.manifest, mapSelection.descriptorPath);
          const composed = composeProjectMap(
            new TextEncoder().encode(workspace.readText(modelRelativePath)),
            mapPath,
            prepared.manifest,
            resolvedMap,
            mapSelection,
          );
          workspace.writeGenerated(mapPath, composed.data);
          warnings.push(...composed.warnings, composed.summary);
        }
        modelRelativePath = mapPath;
        modelPath = workspace.path(modelRelativePath);
      }
      report(0.84, '编译模型与物理数据');
      console.info('[MuJoCo] 编译模型', modelPath);
      nextSession = new SimulationSession(module, modelPath, warnings);
      if (entry?.format === 'urdf' && urdfMode === 'native') {
        const offset = nextSession.alignLowestPointToGround();
        warnings.push(`原生 URDF 已整体平移 ${offset.toFixed(4)} m，使最低点位于 z=0`);
      }
      if (warnings.length) console.info('[MuJoCo] 兼容与地图处理', warnings);
      console.info('[MuJoCo] 模型编译完成');
      report(0.94, '生成初始仿真状态');
      const snapshot = nextSession.snapshot();
      if (this.disposed || generation !== this.loadGeneration) throw new Error('模型加载已取消');
      this.releaseRetired();
      this.retiredSession = this.session;
      this.retiredWorkspace = this.workspace;
      this.retiredSupportFiles = this.supportFiles;
      this.session = nextSession;
      this.workspace = workspace;
      this.supportFiles = supportFiles;
      nextSession = null;
      console.info('[MuJoCo] 状态快照完成');
      return snapshot;
    } catch (error) {
      nextSession?.dispose();
      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);
  }
  setActuatorParameters(id: number, parameters: ActuatorParameters): boolean {
    return this.session?.setActuatorParameters(id, parameters) ?? false;
  }
  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();
  }
  async loadPythonController(source: string, path: string): Promise<ControllerStatus> {
    if (!this.session) throw new Error('请先加载模型');
    return this.session.loadPythonController(source, path);
  }
  setControllerEnabled(enabled: boolean): void {
    this.session?.setControllerEnabled(enabled);
  }
  sendControllerCommand(command: ControllerCommand): void {
    this.session?.sendControllerCommand(command);
  }
  removeController(): void {
    this.session?.removeController();
  }
  async loadRLPolicy(model: Uint8Array, path: string): Promise<RLPolicyStatus> {
    if (!this.session) throw new Error('请先加载模型');
    return this.session.loadRLPolicy(model, path);
  }
  setRLPolicyEnabled(enabled: boolean): void {
    this.session?.setRLPolicyEnabled(enabled);
  }
  setRLCommand(command: RLCommand): void {
    this.session?.setRLCommand(command);
  }
  removeRLPolicy(): void {
    this.session?.removeRLPolicy();
  }
  configureDataRecorder(config: Partial<DataRecorderConfig>): DataRecorderStatus | undefined {
    return this.session?.configureDataRecorder(config);
  }
  startDataRecording(): DataRecorderStatus | undefined {
    return this.session?.startDataRecording();
  }
  stopDataRecording(): DataRecorderStatus | undefined {
    return this.session?.stopDataRecording();
  }
  clearDataRecording(): DataRecorderStatus | undefined {
    return this.session?.clearDataRecording();
  }
  registerDataChannel(channel: TelemetryChannel): (() => void) | undefined {
    return this.session?.registerDataChannel(channel);
  }
  exportDataRecording(format: 'csv' | 'json'): Uint8Array {
    if (!this.session) throw new Error('请先加载模型');
    return this.session.exportDataRecording(format);
  }
  cachedSupportFiles(): ProjectFile[] {
    return this.supportFiles.map((file) => ({ ...file, data: file.data.slice() }));
  }
  releaseRetired(): void {
    this.retiredSession?.dispose();
    this.retiredSession = null;
    this.retiredWorkspace?.dispose();
    this.retiredWorkspace = null;
    this.retiredSupportFiles = [];
  }
  rollbackRetired(): void {
    this.session?.dispose();
    this.workspace?.dispose();
    this.session = this.retiredSession;
    this.workspace = this.retiredWorkspace;
    this.supportFiles = this.retiredSupportFiles;
    this.retiredSession = null;
    this.retiredWorkspace = null;
    this.retiredSupportFiles = [];
  }
  exportMjcf(): Uint8Array {
    if (!this.session || !this.workspace) throw new Error('尚未加载可导出的模型');
    const relative = '.__platform_export__.xml';
    if (this.session.module.mj_saveLastXML(this.workspace.path(relative), this.session.model) === 0)
      throw new Error('MuJoCo 无法生成 MJCF');
    const source = this.workspace.readText(relative),
      document = new DOMParser().parseFromString(source, 'application/xml'),
      actuatorSection = document.querySelector('mujoco > actuator');
    if (document.querySelector('parsererror')) return new TextEncoder().encode(source);
    const snapshot = this.session.snapshot();
    if (actuatorSection) {
      for (const info of snapshot.actuators) {
        const element = Array.from(actuatorSection.children).find(
          (candidate) => candidate.getAttribute('name') === info.name,
        );
        if (!element) continue;
        element.setAttribute('ctrllimited', info.ctrlLimited ? 'true' : 'false');
        element.setAttribute('forcelimited', info.forceLimited ? 'true' : 'false');
      }
    }
    for (const info of snapshot.actuators) {
      if (info.kind !== 'motor' || !info.jointName) continue;
      const joint = Array.from(document.querySelectorAll('worldbody joint[name]')).find(
        (candidate) => candidate.getAttribute('name') === info.jointName,
      );
      if (joint) {
        joint.setAttribute('stiffness', String(info.kp));
        joint.setAttribute('damping', String(info.kv));
      }
    }
    const output = new TextEncoder().encode(new XMLSerializer().serializeToString(document));
    this.workspace.writeGenerated(relative, output);
    return output;
  }
  private releaseCurrent(): void {
    this.releaseRetired();
    this.session?.dispose();
    this.session = null;
    this.workspace?.dispose();
    this.workspace = null;
    this.supportFiles = [];
  }
  dispose(): void {
    this.disposed = true;
    this.loadGeneration += 1;
    this.releaseCurrent();
  }
}