import {
  composeTrainingMap,
  trainingTerrainFromCompiledScene,
  type CompiledMapGeometry,
  type TrainingSceneCoordinates,
} from '../map/trainingMap';
import type { PolicyDeployment, TrainingTerrain } from '../rl/deployment';
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 { composePlacedMapAssets } from '../map/MapStackComposer';
import { DEFAULT_MAP_SELECTION, type MapSelection, type PlacedMapAsset } from '../map/types';

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

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

export interface PhysicsLoadOptions {
  trainingDeployment?: PolicyDeployment;
  /** Candidate policy is initialized/validated before replacing the active session. */
  trainingPolicy?: { data: Uint8Array; path: string };
  urdfMode?: UrdfLoadMode;
  baseMode?: UrdfBaseMode;
  enhancements?: UrdfEnhancementOptions;
  /** @deprecated 单地图兼容入口；新场景应使用 mapAssets。 */
  map?: MapSelection;
  mapAssets?: readonly PlacedMapAsset[];
  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,
    deployment?: PolicyDeployment,
  ): Promise<RLPolicyStatus>;
  setRLPolicyEnabled(enabled: boolean): void;
  setRLCommand(command: RLCommand): void;
  setNavigationTarget(target: [number, number]): void;
  resetNavigationTarget(): 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;
  exportTrainingTerrain(
    assets: readonly PlacedMapAsset[],
    coordinates?: TrainingSceneCoordinates,
  ): TrainingTerrain;
  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 trainingScenes = new WeakMap<
    SimulationSession,
    {
      assets: string;
      geometries: CompiledMapGeometry[];
      initialPose: number[];
      extent: number;
    }
  >();
  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 placedMapAssets = options.mapAssets;
    const hasMaps = placedMapAssets ? placedMapAssets.length > 0 : mapSelection.kind !== 'none';
    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 (hasMaps && !options.trainingDeployment?.terrain) {
        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`;
        const source = new TextEncoder().encode(workspace.readText(modelRelativePath));
        if (placedMapAssets) {
          const composed = composePlacedMapAssets(
            source,
            mapPath,
            prepared.manifest,
            placedMapAssets,
          );
          workspace.writeGenerated(mapPath, composed.data);
          warnings.push(...composed.warnings);
        } else if (mapSelection.kind === 'builtin') {
          const composed = composePhysicalMap(source, mapSelection.config);
          workspace.writeGenerated(mapPath, composed.data);
          if (composed.summary) warnings.push(composed.summary);
        } else if (mapSelection.kind === 'project') {
          const resolvedMap = resolveProjectMap(prepared.manifest, mapSelection.descriptorPath);
          const composed = composeProjectMap(
            source,
            mapPath,
            prepared.manifest,
            resolvedMap,
            mapSelection,
          );
          workspace.writeGenerated(mapPath, composed.data);
          warnings.push(...composed.warnings, composed.summary);
        }
        modelRelativePath = mapPath;
        modelPath = workspace.path(modelRelativePath);
      }
      if (options.trainingDeployment?.terrain) {
        if (entry?.format === 'urdf' && urdfMode === 'native')
          throw new Error('训练地图需要URDF转换模式');
        const intermediate = new SimulationSession(module, modelPath);
        try {
          const flatPath = `${modelRelativePath}.training.xml`;
          if (!module.mj_saveLastXML(workspace.path(flatPath), intermediate.model))
            throw new Error('无法展开训练场景');
          workspace.writeGenerated(
            flatPath,
            composeTrainingMap(
              new TextEncoder().encode(workspace.readText(flatPath)),
              options.trainingDeployment,
            ),
          );
          modelPath = workspace.path(flatPath);
        } finally {
          intermediate.dispose();
        }
        warnings.push('已用策略配套训练布局替换场景地形；编辑器地图未更改。Go2-W动力学不等同Go2。');
      }
      report(0.84, '编译模型与物理数据');
      console.info('[MuJoCo] 编译模型', modelPath);
      nextSession = new SimulationSession(module, modelPath, warnings);
      if (options.trainingDeployment) nextSession.configureDeployment(options.trainingDeployment);
      if (entry?.format === 'urdf' && urdfMode === 'native') {
        const offset = nextSession.alignLowestPointToGround();
        warnings.push(`原生 URDF 已整体平移 ${offset.toFixed(4)} m，使最低点位于 z=0`);
      }
      if (placedMapAssets?.length && !options.trainingDeployment) {
        const sanitize = (id: string) => id.replace(/[^a-zA-Z0-9_-]/g, '_');
        const prefixes = placedMapAssets.map((asset) =>
          asset.selection.kind === 'builtin'
            ? `__platform_map_${sanitize(asset.id)}__`
            : `__platform_map_${sanitize(`${resolveProjectMap(prepared.manifest, asset.selection.descriptorPath).definition.id}_${asset.id}`)}_`,
        );
        const model = nextSession.model,
          data = nextSession.data;
        const geometries: CompiledMapGeometry[] = [];
        for (let id = 0; id < model.ngeom; id++) {
          const geom = model.geom(id);
          try {
            const name = geom.name;
            geometries.push({
              name,
              type: Number(model.geom_type[id]),
              position: Array.from(data.geom_xpos.slice(id * 3, id * 3 + 3)),
              rotation: Array.from(data.geom_xmat.slice(id * 9, id * 9 + 9)),
              size: Array.from(model.geom_size.slice(id * 3, id * 3 + 3)),
              friction: Array.from(model.geom_friction.slice(id * 3, id * 3 + 3)),
              collision: Boolean(model.geom_contype[id] || model.geom_conaffinity[id]),
              static: Number(model.body_weldid[Number(model.geom_bodyid[id])]) === 0,
              map:
                prefixes.some((prefix) => name.startsWith(prefix)) ||
                name === '__platform_map_ground__',
            });
          } finally {
            geom.delete();
          }
        }
        const roots = Array.from({ length: model.njnt }, (_, id) => id).filter(
          (id) => Number(model.jnt_type[id]) === 0,
        );
        if (roots.length === 1) {
          const address = Number(model.jnt_qposadr[roots[0]]);
          this.trainingScenes.set(nextSession, {
            assets: JSON.stringify(placedMapAssets),
            geometries,
            initialPose: Array.from(data.qpos.slice(address, address + 7)),
            extent: Math.max(
              4,
              ...placedMapAssets.map((asset) =>
                asset.selection.kind === 'builtin'
                  ? Math.max(
                      Math.abs(asset.selection.config.positionX),
                      Math.abs(asset.selection.config.positionY),
                    ) +
                    asset.selection.config.size / 2
                  : 0,
              ),
            ),
          });
        }
      }
      if (options.trainingPolicy) {
        if (!options.trainingDeployment?.terrain) throw new Error('事务策略加载需要配套训练地图');
        report(0.9, '校验候选场景的 ONNX 策略');
        await nextSession.loadRLPolicy(
          options.trainingPolicy.data,
          options.trainingPolicy.path,
          options.trainingDeployment,
        );
        // Configure/reset already supplied the correct spawn. Enable while still paused;
        // the candidate cannot advance until the viewer transaction commits.
        nextSession.setRLPolicyEnabled(true);
      }
      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 },
      );
    }
  }
  exportTrainingTerrain(
    assets: readonly PlacedMapAsset[],
    coordinates?: TrainingSceneCoordinates,
  ): TrainingTerrain {
    const scene = this.session && this.trainingScenes.get(this.session);
    if (!scene || scene.assets !== JSON.stringify(assets))
      throw new Error('没有匹配的已应用碰撞场景/唯一浮动机器人，或场景已过时；请先应用地图');
    return trainingTerrainFromCompiledScene(
      scene.geometries,
      scene.initialPose,
      scene.extent,
      coordinates,
    );
  }
  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,
    deployment?: PolicyDeployment,
  ): Promise<RLPolicyStatus> {
    if (!this.session) throw new Error('请先加载模型');
    return this.session.loadRLPolicy(model, path, deployment);
  }
  setRLPolicyEnabled(enabled: boolean): void {
    this.session?.setRLPolicyEnabled(enabled);
  }
  setNavigationTarget(target: [number, number]): void {
    this.session?.setNavigationTarget(target);
  }
  resetNavigationTarget(): void {
    this.session?.resetNavigationTarget();
  }
  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 {
    const retiredSession = this.retiredSession;
    const retiredWorkspace = this.retiredWorkspace;
    this.retiredSession = null;
    this.retiredWorkspace = null;
    this.retiredSupportFiles = [];
    try {
      retiredSession?.dispose();
    } catch (error) {
      console.warn('[MuJoCo] 释放旧仿真会话失败', error);
    }
    try {
      retiredWorkspace?.dispose();
    } catch (error) {
      console.warn('[MuJoCo] 释放旧内存工作区失败', error);
    }
  }
  rollbackRetired(): void {
    const failedSession = this.session;
    const failedWorkspace = this.workspace;
    this.session = this.retiredSession;
    this.workspace = this.retiredWorkspace;
    this.supportFiles = this.retiredSupportFiles;
    this.retiredSession = null;
    this.retiredWorkspace = null;
    this.retiredSupportFiles = [];
    try {
      failedSession?.dispose();
    } catch (error) {
      console.warn('[MuJoCo] 释放失败的候选仿真会话失败', error);
    }
    try {
      failedWorkspace?.dispose();
    } catch (error) {
      console.warn('[MuJoCo] 释放失败的候选内存工作区失败', error);
    }
  }
  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();
  }
}