Mujoco_WASM/wasm/web_platform/src/project/importer.ts

import {unzipSync} from 'fflate';
import {DEFAULT_IMPORT_LIMITS, type ImportLimits, type ModelEntry, type ProjectFile, type ProjectManifest} from './types';
import {convertDaeToObj} from './daeConverter';

const decoder = new TextDecoder('utf-8', {fatal: false});

export class ProjectImportError extends Error {
  constructor(message: string, readonly path?: string) { super(message); this.name = 'ProjectImportError'; }
}

export function normalizeProjectPath(input: string): string {
  const path = input.replaceAll('\\', '/').replace(/^\.\//, '');
  if (!path || path.startsWith('/') || path.includes('\0') || /^[A-Za-z]:/.test(path)) throw new ProjectImportError('不允许绝对路径或空路径', input);
  const parts = path.split('/').filter((part) => part !== '' && part !== '.');
  if (!parts.length || parts.some((part) => part === '..')) throw new ProjectImportError('路径包含越界片段', input);
  return parts.join('/');
}

function checkEncryptedZip(data: Uint8Array): void {
  for (let i = 0; i + 8 < data.length; i++) {
    if (data[i] === 0x50 && data[i + 1] === 0x4b && (data[i + 2] === 0x03 || data[i + 2] === 0x01) && (data[i + 3] === 0x04 || data[i + 3] === 0x02)) {
      const flags = data[i + 6] | (data[i + 7] << 8);
      if ((flags & 1) !== 0) throw new ProjectImportError('不支持加密 ZIP');
    }
  }
}

function enforceLimits(files: ProjectFile[], limits: ImportLimits): void {
  if (files.length > limits.maxFiles) throw new ProjectImportError(`文件数量超过限制（${limits.maxFiles}）`);
  let total = 0;
  const seen = new Set<string>();
  for (const file of files) {
    if (seen.has(file.path)) throw new ProjectImportError('工程中存在同名路径', file.path);
    seen.add(file.path);
    if (file.size > limits.maxFileBytes) throw new ProjectImportError(`单文件超过限制（${limits.maxFileBytes} 字节）`, file.path);
    total += file.size;
    if (total > limits.maxTotalBytes) throw new ProjectImportError(`工程总大小超过限制（${limits.maxTotalBytes} 字节）`);
  }
}

export function discoverEntries(files: ProjectFile[]): ModelEntry[] {
  return files.flatMap((file): ModelEntry[] => {
    if (!/\.(xml|urdf)$/i.test(file.path)) return [];
    const head = decoder.decode(file.data.subarray(0, Math.min(file.data.length, 256 * 1024))).replace(/^\uFEFF/, '');
    const format = /<robot(?:\s|\/?>)/i.test(head) ? 'urdf' : /<mujoco(?:\s|\/?>)/i.test(head) ? 'mjcf' : undefined;
    return format ? [{path: file.path, format, label: `${file.path} (${format.toUpperCase()})`}] : [];
  });
}

export interface PreparedProject {
  manifest: ProjectManifest;
  warnings: string[];
}

function relativeProjectPath(fromFile: string, toFile: string): string {
  const from = fromFile.split('/').slice(0, -1);
  const to = toFile.split('/');
  while (from.length && to.length && from[0] === to[0]) { from.shift(); to.shift(); }
  return `${'../'.repeat(from.length)}${to.join('/')}` || './';
}

function resolveProjectReference(fromFile:string,reference:string):string|undefined {
  if(/^[a-z][a-z\d+.-]*:/i.test(reference))return;
  let decoded:string;
  try{decoded=decodeURIComponent(reference.split(/[?#]/,1)[0]);}catch{return;}
  const parts=fromFile.split('/').slice(0,-1);
  for(const part of decoded.replaceAll('\\','/').split('/')){
    if(!part||part==='.')continue;
    if(part==='..'){if(!parts.length)return;parts.pop();}
    else parts.push(part);
  }
  return parts.join('/');
}

function generatedObjPath(daePath:string,occupied:Set<string>):string {
  const base=daePath.replace(/\.dae$/i,'');
  let candidate=`${base}.mujoco.obj`;
  for(let index=2;occupied.has(candidate);index+=1)candidate=`${base}.mujoco-${index}.obj`;
  occupied.add(candidate);
  return candidate;
}

/** Normalizes common ROS URDF constructs that MuJoCo's stricter parser rejects. */
export function prepareProjectForMujoco(manifest: ProjectManifest, entryPath: string): PreparedProject {
  const entry = manifest.entries.find((candidate) => candidate.path === entryPath);
  if (entry?.format !== 'urdf') return {manifest, warnings: []};
  const source = manifest.files.find((file) => file.path === entryPath);
  if (!source) return {manifest, warnings: []};
  const document = new DOMParser().parseFromString(decoder.decode(source.data), 'application/xml');
  if (document.querySelector('parsererror')) return {manifest, warnings: []};

  const warnings: string[] = [];
  const robot=document.documentElement;
  let mujoco=Array.from(robot.children).find(child=>child.tagName==='mujoco');
  if(!mujoco){mujoco=document.createElement('mujoco');robot.prepend(mujoco);}
  let compiler=Array.from(mujoco.children).find(child=>child.tagName==='compiler');
  if(!compiler){compiler=document.createElement('compiler');mujoco.append(compiler);}
  compiler.setAttribute('discardvisual','false');
  compiler.setAttribute('fusestatic','false');

  let removedMaterials = 0;
  for (const visual of Array.from(document.querySelectorAll('visual'))) {
    const materials = Array.from(visual.children).filter((child) => child.tagName === 'material');
    for (const duplicate of materials.slice(1)) { duplicate.remove(); removedMaterials += 1; }
  }
  if (removedMaterials) warnings.push(`为兼容 MuJoCo，已移除 visual 中 ${removedMaterials} 个重复 material（保留第一个）`);

  const paths = manifest.files.map((file) => file.path);
  let rewrittenUris = 0;
  const unresolved: string[] = [];
  for (const element of Array.from(document.querySelectorAll('[filename]'))) {
    const value = element.getAttribute('filename');
    if (!value?.startsWith('package://')) continue;
    const packagePath = normalizeProjectPath(value.slice('package://'.length));
    const target = paths.find((path) => path === packagePath) ?? paths.find((path) => path.endsWith(`/${packagePath}`));
    if (!target) { unresolved.push(value); continue; }
    element.setAttribute('filename', relativeProjectPath(entryPath, target));
    rewrittenUris += 1;
  }
  if (rewrittenUris) warnings.push(`已将 ${rewrittenUris} 个 package:// 资源地址改写为工程内相对路径`);
  if (unresolved.length) warnings.push(`有 ${unresolved.length} 个 package:// 资源未在工程中找到`);

  const occupied=new Set(manifest.files.map(file=>file.path));
  const converted=new Map<string,ProjectFile>();
  let convertedDaeReferences=0;
  let removedDaeVisuals=0;
  let daeCollisionFallbacks=0;
  for(const mesh of Array.from(document.querySelectorAll('mesh[filename]'))){
    const filename=mesh.getAttribute('filename');
    if(!filename?.toLowerCase().split(/[?#]/)[0].endsWith('.dae'))continue;
    const daePath=resolveProjectReference(entryPath,filename);
    const daeFile=daePath?manifest.files.find(file=>file.path===daePath):undefined;
    try{
      if(!daeFile||!daePath)throw new Error('工程中找不到 DAE 文件');
      let objFile=converted.get(daePath);
      if(!objFile){
        const data=convertDaeToObj(daeFile.data,daePath);
        if(data.byteLength>DEFAULT_IMPORT_LIMITS.maxFileBytes)throw new Error('转换后的 OBJ 超过单文件大小限制');
        objFile={path:generatedObjPath(daePath,occupied),data,size:data.byteLength,source:daeFile.source,mimeType:'text/plain'};
        converted.set(daePath,objFile);
      }
      mesh.setAttribute('filename',relativeProjectPath(entryPath,objFile.path));
      convertedDaeReferences+=1;
    }catch(error){
      console.warn(`[MuJoCo] DAE 转换失败：${filename}`,error);
      const visual=mesh.closest('visual');
      if(visual){visual.remove();removedDaeVisuals+=1;}
      else if(mesh.closest('collision')){
        const sphere=document.createElement('sphere');sphere.setAttribute('radius','0.05');mesh.replaceWith(sphere);daeCollisionFallbacks+=1;
      }
    }
  }
  if(convertedDaeReferences)warnings.push(`已将 ${converted.size} 个 DAE 文件转换为 OBJ，供 ${convertedDaeReferences} 个 visual/collision 使用`);
  if(removedDaeVisuals)warnings.push(`${removedDaeVisuals} 个 DAE visual 转换失败，已移除并使用其他 collision 几何显示/仿真`);
  if(daeCollisionFallbacks)warnings.push(`${daeCollisionFallbacks} 个 DAE collision 转换失败，已替换为半径 0.05 m 的占位球体；碰撞精度会降低`);

  const xml=new TextEncoder().encode(new XMLSerializer().serializeToString(document));
  const replacement:ProjectFile={...source,data:xml,size:xml.byteLength};
  const generated=Array.from(converted.values());
  const files=[...manifest.files.map(file=>file===source?replacement:file),...generated];
  return {manifest:{...manifest,files,totalBytes:files.reduce((total,file)=>total+file.size,0)},warnings};
}

export function choosePreferredEntry(entries: ModelEntry[]): string | undefined {
  if (entries.length === 1) return entries[0].path;
  const rootPreferred = entries.find((e) => !e.path.includes('/') && /^(model|scene)\.xml$/i.test(e.path));
  if (rootPreferred) return rootPreferred.path;
  const urdfs = entries.filter((e) => e.format === 'urdf');
  return urdfs.length === 1 ? urdfs[0].path : undefined;
}

function manifest(name: string, files: ProjectFile[]): ProjectManifest {
  const entries = discoverEntries(files);
  if (!entries.length) throw new ProjectImportError('未发现包含 <mujoco> 或 <robot> 根元素的 XML/URDF 入口');
  return {id: `${Date.now().toString(36)}-${crypto.randomUUID().slice(0, 8)}`, name, files, entries, selectedEntry: choosePreferredEntry(entries), totalBytes: files.reduce((n, f) => n + f.size, 0)};
}

export async function importBrowserFiles(input: File[], limits: ImportLimits = DEFAULT_IMPORT_LIMITS): Promise<ProjectManifest> {
  if (!input.length) throw new ProjectImportError('未选择文件');
  if (input.length === 1 && /\.zip$/i.test(input[0].name)) {
    if (input[0].size > limits.maxZipBytes) throw new ProjectImportError(`ZIP 超过限制（${limits.maxZipBytes} 字节）`);
    const bytes = new Uint8Array(await input[0].arrayBuffer()); checkEncryptedZip(bytes);
    let unpacked: Record<string, Uint8Array>;
    try {
      let fileCount = 0;
      let expandedBytes = 0;
      unpacked = unzipSync(bytes, {filter: (entry) => {
        if (entry.name.endsWith('/')) return false;
        normalizeProjectPath(entry.name);
        fileCount += 1;
        expandedBytes += entry.originalSize;
        if (fileCount > limits.maxFiles) throw new ProjectImportError(`文件数量超过限制（${limits.maxFiles}）`);
        if (entry.originalSize > limits.maxFileBytes) throw new ProjectImportError(`单文件超过限制（${limits.maxFileBytes} 字节）`, entry.name);
        if (expandedBytes > limits.maxTotalBytes) throw new ProjectImportError(`ZIP 解压后总大小超过限制（${limits.maxTotalBytes} 字节）`);
        return true;
      }});
    } catch (error) {
      if (error instanceof ProjectImportError) throw error;
      throw new ProjectImportError(`ZIP 解压失败：${error instanceof Error ? error.message : String(error)}`);
    }
    const files = Object.entries(unpacked).filter(([path]) => !path.endsWith('/')).map(([path, data]): ProjectFile => ({path: normalizeProjectPath(path), data, size: data.byteLength, source: 'zip', mimeType: ''}));
    enforceLimits(files, limits); return manifest(input[0].name.replace(/\.zip$/i, ''), files);
  }
  const files = await Promise.all(input.map(async (file): Promise<ProjectFile> => {
    const relative = (file as File & {webkitRelativePath?: string}).webkitRelativePath || file.name;
    const data = new Uint8Array(await file.arrayBuffer());
    return {path: normalizeProjectPath(relative), data, size: data.byteLength, source: relative === file.name ? 'file' : 'directory', mimeType: file.type};
  }));
  enforceLimits(files, limits); return manifest(files[0].path.split('/')[0] || '工程', files);
}

interface LegacyEntry {isFile: boolean; isDirectory: boolean; name: string; file(cb: (file: File) => void, err: (e: DOMException) => void): void; createReader(): {readEntries(cb: (entries: LegacyEntry[]) => void, err: (e: DOMException) => void): void};}
async function readEntry(entry: LegacyEntry, prefix = ''): Promise<File[]> {
  if (entry.isFile) return [await new Promise<File>((resolve, reject) => entry.file((file) => {Object.defineProperty(file, 'webkitRelativePath', {value: `${prefix}${file.name}`}); resolve(file);}, reject))];
  const reader = entry.createReader(); const children: LegacyEntry[] = [];
  for (;;) { const batch = await new Promise<LegacyEntry[]>((resolve, reject) => reader.readEntries(resolve, reject)); if (!batch.length) break; children.push(...batch); }
  return (await Promise.all(children.map((child) => readEntry(child, `${prefix}${entry.name}/`)))).flat();
}

export async function filesFromDrop(items: DataTransferItemList, fallback: FileList): Promise<File[]> {
  const entries = Array.from(items).map((item) => (item as unknown as {webkitGetAsEntry?: () => LegacyEntry | null}).webkitGetAsEntry?.() ?? null).filter((entry): entry is LegacyEntry => entry !== null);
  return entries.length ? (await Promise.all(entries.map((entry) => readEntry(entry)))).flat() : Array.from(fallback);
}