PRA-only-demo/src/main/local-vision/ui-detect-service.ts

import fs from "node:fs";
import path from "node:path";
import { Worker } from "node:worker_threads";
import { app } from "electron";
import { RuntimeError } from "../engine/types";
import type { LocalUiBox } from "./ui-detect-worker";

export interface UiDetectOptions {
  image: Buffer;
  width: number;
  height: number;
  modelPath?: string;
  scoreThreshold?: number;
  iouThreshold?: number;
}

interface DetectRequest {
  id: number;
  image: Buffer;
  width: number;
  height: number;
  scoreThreshold?: number;
  iouThreshold?: number;
}

interface DetectResponse {
  id: number;
  boxes?: LocalUiBox[];
  error?: string;
  elapsedMs?: number;
}

const workers = new Map<string, Worker>();
const pending = new Map<number, { resolve: (boxes: LocalUiBox[]) => void; reject: (error: unknown) => void; worker: Worker }>();
let nextRequestId = 1;
const MAX_LOADED_WORKERS = 4;

export function localUiModelPath(): string {
  const candidates = [
    path.join(process.resourcesPath || "", "models", "ui-icon-detect", "model.onnx"),
    path.join(app.getAppPath(), "resources", "models", "ui-icon-detect", "model.onnx")
  ].filter(Boolean);
  const found = candidates.find(candidate => fs.existsSync(candidate));
  if (!found) {
    throw new RuntimeError("local_vision_model_missing", "本地 UI 检测模型缺失", {
      recommendation: "重新安装应用，或补齐 resources/models/ui-icon-detect/model.onnx"
    });
  }
  return found;
}

export function defaultUiModelPath(): string {
  return localUiModelPath();
}

function resolveModelPath(candidate?: string): string {
  if (!candidate) return localUiModelPath();
  const normalized = path.resolve(candidate);
  if (!fs.existsSync(normalized) || !fs.statSync(normalized).isFile()) {
    throw new RuntimeError("local_vision_model_missing", "模型文件不存在：" + normalized);
  }
  if (path.extname(normalized).toLowerCase() !== ".onnx") {
    throw new RuntimeError("local_vision_model_invalid", "请选择 ONNX 模型文件");
  }
  return normalized;
}

function createWorker(modelPath: string): Worker {
  const next = new Worker(path.join(__dirname, "local-vision/ui-detect-worker.js"), {
    workerData: { modelPath },
    resourceLimits: { maxOldGenerationSizeMb: 512 }
  });
  next.on("message", (message: DetectResponse) => {
    const entry = pending.get(message.id);
    if (!entry || entry.worker !== next) return;
    pending.delete(message.id);
    if (message.error) {
      entry.reject(new RuntimeError("local_vision_failed", "本地 UI 检测失败：" + message.error));
      return;
    }
    entry.resolve(message.boxes || []);
  });
  next.on("error", error => {
    for (const [id, entry] of pending) {
      if (entry.worker === next) {
        entry.reject(error);
        pending.delete(id);
      }
    }
    workers.delete(modelPath);
  });
  next.on("exit", code => {
    for (const [id, entry] of pending) {
      if (entry.worker === next) {
        entry.reject(new RuntimeError("local_vision_failed", "本地 UI 检测进程已退出：" + code));
        pending.delete(id);
      }
    }
    if (workers.get(modelPath) === next) workers.delete(modelPath);
  });
  return next;
}

function getWorker(modelPath: string): Worker {
  const existing = workers.get(modelPath);
  if (existing) return existing;
  while (workers.size >= MAX_LOADED_WORKERS) {
    const oldest = workers.keys().next().value;
    if (!oldest) break;
    const worker = workers.get(oldest);
    workers.delete(oldest);
    void worker?.terminate();
  }
  const created = createWorker(modelPath);
  workers.set(modelPath, created);
  return created;
}

export async function detectUiElements(input: UiDetectOptions): Promise<LocalUiBox[]> {
  const modelPath = resolveModelPath(input.modelPath);
  const current = getWorker(modelPath);
  const nextId = nextRequestId++;

  const request: DetectRequest = {
    id: nextId,
    image: input.image,
    width: input.width,
    height: input.height,
    scoreThreshold: input.scoreThreshold,
    iouThreshold: input.iouThreshold
  };
  return new Promise<LocalUiBox[]>((resolve, reject) => {
    pending.set(nextId, { resolve, reject, worker: current });
    current.postMessage(request);
  });
}

export function closeLocalVisionWorkers(): void {
  for (const [modelPath, worker] of workers) {
    workers.delete(modelPath);
    void worker.terminate();
  }
  pending.clear();
}