massive-cloud/server/deviationHeatmapJob.ts

import * as fs from 'fs';
import { fork } from 'node:child_process';
import { fileURLToPath } from 'node:url';
import path from 'node:path';
import * as db from './db.js';
import {
  computeDeviationHeatmapFromObj,
  computeDeviationHeatmapFromStl,
} from './cadHeatmapCompute.js';
import { pickCadFileForHeatmap } from './cadFilePick.js';
import { sampleTaskPointclouds } from './sampleTaskPointclouds.js';
import { deviationGpuEnabled, tryRunDeviationGpu } from './deviationGpuCompute.js';

const DATA_DIR = process.env.DATA_DIR || '/data';
const TASKS_DIR = path.join(DATA_DIR, 'tasks');

async function reportProgress(resultId: number, progress: number, stage: string): Promise<void> {
  const p = Math.min(1, Math.max(0, progress));
  await db.updateTaskAnalysisResult(resultId, {
    resultData: {
      kind: 'obj_deviation',
      progress: p,
      stage,
    },
  });
}

function resolveRunnerEntry(): string {
  const dir = path.dirname(fileURLToPath(import.meta.url));
  if (process.env.NODE_ENV === 'production') {
    return path.join(dir, 'deviationHeatmapJobRunner.js');
  }
  return path.join(dir, 'deviationHeatmapJobRunner.ts');
}

export function enqueueDeviationHeatmapJob(resultId: number, taskId: number, totalPoints: number): void {
  const script = resolveRunnerEntry();
  const isProd = process.env.NODE_ENV === 'production';
  const child = fork(script, [String(resultId), String(taskId), String(totalPoints)], {
    env: { ...process.env },
    stdio: ['ignore', 'inherit', 'inherit', 'ipc'],
    execArgv: isProd ? [] : ['--import', 'tsx'],
  });
  child.on('error', (err) => {
    console.error('[DeviationHeatmapJob] fork 失败:', err);
  });
  child.on('exit', (code, signal) => {
    if (code !== 0 && code != null) {
      console.error(`[DeviationHeatmapJob] 子进程退出 code=${code} signal=${signal ?? ''}`);
    }
  });
}

export async function runDeviationHeatmapJob(resultId: number, taskId: number, totalPoints: number): Promise<void> {
  try {
    await reportProgress(resultId, 0.06, '加载任务与 CAD…');
    const task = await db.getInspectionTaskById(taskId);
    if (!task) throw new Error('task not found');

    const cadFiles = await db.listModelCadFiles(task.modelId);
    const picked = pickCadFileForHeatmap(cadFiles);
    if (!picked) {
      throw new Error(
        '当前车型下未找到 STL 或 OBJ 的 CAD。请在与该检测任务相同的车型下上传 CAD，或确认文件扩展名为 .stl / .obj',
      );
    }
    if (!fs.existsSync(picked.file.filePath)) throw new Error('CAD 文件不存在');

    await reportProgress(resultId, 0.14, '点云采样中…');
    const allPoints = await sampleTaskPointclouds(taskId, totalPoints);
    if (allPoints.length === 0) throw new Error('任务下无可用点云采样');

    // ── GPU 路径 ──────────────────────────────────────────────────────
    let computed: { colors: Float32Array; vertexCount: number; dispMin: number; dispMax: number } | null = null;

    if (deviationGpuEnabled()) {
      await reportProgress(resultId, 0.42, '尝试 GPU 加速计算偏差…');
      console.log('[DeviationHeatmapJob] GPU enabled, trying CUDA path…');
      const gpuResult = tryRunDeviationGpu(picked.file.filePath, allPoints);
      if (gpuResult) {
        console.log(`[DeviationHeatmapJob] GPU 完成: verts=${gpuResult.vertexCount}`);
        computed = gpuResult;
      } else {
        console.warn('[DeviationHeatmapJob] GPU 不可用，回退 CPU');
      }
    }

    // ── CPU 回退 ──────────────────────────────────────────────────────
    if (!computed) {
      await reportProgress(resultId, 0.42, '计算偏差并着色（CPU，耗时较长）…');
      const cadBuf = fs.readFileSync(picked.file.filePath);
      computed =
        picked.ext === 'stl'
          ? computeDeviationHeatmapFromStl(cadBuf, allPoints)
          : computeDeviationHeatmapFromObj(cadBuf, allPoints);
    }
    if (!computed) throw new Error('偏差热力计算失败');

    await reportProgress(resultId, 0.88, '写入缓存文件…');
    const outDir = path.join(TASKS_DIR, String(taskId), 'analysis');
    fs.mkdirSync(outDir, { recursive: true });
    const binName = `deviation_${resultId}.bin`;
    const binPath = path.join(outDir, binName);
    fs.writeFileSync(
      binPath,
      Buffer.from(computed.colors.buffer, computed.colors.byteOffset, computed.colors.byteLength),
    );

    await db.updateTaskAnalysisResult(resultId, {
      status: 'completed',
      resultData: {
        kind: 'obj_deviation',
        vertexCount: computed.vertexCount,
        dispMin: computed.dispMin,
        dispMax: computed.dispMax,
        binName,
      },
    });
    console.log(
      `[DeviationHeatmapJob] ok resultId=${resultId} taskId=${taskId} verts=${computed.vertexCount}`,
    );
  } catch (err) {
    console.error('[DeviationHeatmapJob]', err);
    await db.updateTaskAnalysisResult(resultId, {
      status: 'error',
      resultData: {
        kind: 'obj_deviation',
        error: err instanceof Error ? err.message : String(err),
      },
    });
  }
}