import fs from "node:fs";
import { parentPort, workerData } from "node:worker_threads";
import { InferenceSession, Tensor } from "onnxruntime-node";
import { createCanvas, loadImage } from "@napi-rs/canvas";

export interface LocalUiBox {
  x: number;
  y: number;
  width: number;
  height: number;
  category: string;
  confidence: 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;
}

type ModelOutput = { output0: Tensor };
type TensorLike = { dims: number[]; data: Float32Array };

const inputName = "images";
const outputName = "output0";
const inputSize = 640;
const defaultScoreThreshold = 0.25;
const defaultIouThreshold = 0.42;
const maxBoxes = 160;

const startupModelPath = String((workerData as { modelPath?: string }).modelPath || "");
const modelLabels = readModelLabels(startupModelPath);

let sessionPromise: Promise<InferenceSession> | null = null;

function loadSession(): Promise<InferenceSession> {
  if (!sessionPromise) {
    const modelPath = String((workerData as { modelPath?: string }).modelPath || "");
    sessionPromise = InferenceSession.create(modelPath, {
      executionProviders: ["cpu"],
      graphOptimizationLevel: "all",
      executionMode: "sequential"
    });
  }
  return sessionPromise;
}

function iou(a: LocalUiBox, b: LocalUiBox): number {
  const x1 = Math.max(a.x, b.x);
  const y1 = Math.max(a.y, b.y);
  const x2 = Math.min(a.x + a.width, b.x + b.width);
  const y2 = Math.min(a.y + a.height, b.y + b.height);
  const intersection = Math.max(0, x2 - x1) * Math.max(0, y2 - y1);
  if (intersection <= 0) return 0;
  return intersection / (a.width * a.height + b.width * b.height - intersection);
}

function nms(boxes: LocalUiBox[], iouThreshold: number): LocalUiBox[] {
  const sorted = [...boxes].sort((a, b) => b.confidence - a.confidence);
  const kept: LocalUiBox[] = [];
  for (const box of sorted) {
    if (kept.every(item => iou(item, box) <= iouThreshold)) kept.push(box);
  }
  return kept;
}

function tensorAccessor(output: TensorLike, anchorCount: number, channelCount: number, channelFirst: boolean) {
  return (channel: number, anchor: number): number => channelFirst
    ? output.data[channel * anchorCount + anchor]
    : output.data[anchor * channelCount + channel];
}

// YOLO 导出的 ONNX 通常不含标签元数据。约定同名 *.labels.json 保存类别：
// 1) ["button", "input"] 2) { "labels": ["button", "input"] }
function readModelLabels(modelPath: string): string[] | undefined {
  const sidecarPath = modelPath.replace(/\.onnx$/i, ".labels.json");
  try {
    const parsed: unknown = JSON.parse(fs.readFileSync(sidecarPath, "utf8"));
    const values = Array.isArray(parsed) ? parsed : (parsed as { labels?: unknown })?.labels;
    if (!Array.isArray(values)) return undefined;
    const labels = values
      .filter((item): item is string | number => typeof item === "string" || typeof item === "number")
      .map(item => String(item))
      .filter(Boolean);
    return labels.length ? labels : undefined;
  } catch {
    return undefined;
  }
}

function extractBoxes(
  output: TensorLike,
  sourceWidth: number,
  sourceHeight: number,
  scoreThreshold: number,
  iouThreshold: number
): LocalUiBox[] {
  // YOLO/OmniParser exports are commonly [1,C,N] or [1,N,C]. C=4+classCount.
  if (output.dims.length < 2) throw new Error(`不支持的模型输出维度：${output.dims.join("x")}`);
  const [, first, second] = output.dims;
  const channelFirst = first <= second;
  const channelCount = channelFirst ? first : second;
  const anchorCount = channelFirst ? second : first;
  if (channelCount < 5) throw new Error(`模型输出通道数异常：${channelCount}`);
  const at = tensorAccessor(output, anchorCount, channelCount, channelFirst);
  const classCount = channelCount - 4;

  const padded = Math.max(sourceWidth, sourceHeight);
  const scale = inputSize / padded;
  const drawnWidth = Math.round(sourceWidth * scale);
  const drawnHeight = Math.round(sourceHeight * scale);
  const offsetX = (inputSize - drawnWidth) / 2;
  const offsetY = (inputSize - drawnHeight) / 2;
  const minSide = Math.max(8, Math.min(sourceWidth, sourceHeight) * 0.006);
  const labels = modelLabels;
  const raw: LocalUiBox[] = [];

  for (let index = 0; index < anchorCount; index++) {
    let confidence = 0;
    let classIndex = 0;
    if (classCount === 1) {
      confidence = at(4, index);
    } else {
      for (let candidate = 0; candidate < classCount; candidate++) {
        const score = at(4 + candidate, index);
        if (score > confidence) {
          confidence = score;
          classIndex = candidate;
        }
      }
    }
    if (!Number.isFinite(confidence) || confidence < scoreThreshold) continue;
    const cx = at(0, index);
    const cy = at(1, index);
    const width = at(2, index);
    const height = at(3, index);
    let x = (cx - width / 2 - offsetX) / scale;
    let y = (cy - height / 2 - offsetY) / scale;
    let boxWidth = width / scale;
    let boxHeight = height / scale;
    x = Math.max(0, Math.min(sourceWidth - 1, x));
    y = Math.max(0, Math.min(sourceHeight - 1, y));
    boxWidth = Math.max(2, Math.min(sourceWidth - x, boxWidth));
    boxHeight = Math.max(2, Math.min(sourceHeight - y, boxHeight));
    if (boxWidth < minSide || boxHeight < minSide) continue;
    const areaRatio = (boxWidth * boxHeight) / (sourceWidth * sourceHeight);
    if (areaRatio > 0.92) continue;
    raw.push({
      x: Math.round(x),
      y: Math.round(y),
      width: Math.round(boxWidth),
      height: Math.round(boxHeight),
      category: labels?.[classIndex]
        ?? (classCount <= 1 ? "控件" : `类别${classIndex + 1}`),
      confidence: Math.round(Math.max(0, Math.min(1, confidence)) * 1000) / 1000
    });
  }

  return nms(raw, iouThreshold)
    .sort((a, b) => b.confidence - a.confidence || a.y - b.y || a.x - b.x)
    .slice(0, maxBoxes)
    .sort((a, b) => a.y - b.y || a.x - b.x);
}

async function detect(request: DetectRequest): Promise<LocalUiBox[]> {
  const model = await loadSession();
  const source = await loadImage(request.image);
  const width = request.width || source.width;
  const height = request.height || source.height;

  const canvas = createCanvas(inputSize, inputSize);
  const context = canvas.getContext("2d");
  context.fillStyle = "#727272";
  context.fillRect(0, 0, inputSize, inputSize);
  context.imageSmoothingEnabled = true;
  context.imageSmoothingQuality = "high";
  const scale = Math.min(inputSize / width, inputSize / height);
  context.drawImage(
    source,
    (inputSize - width * scale) / 2,
    (inputSize - height * scale) / 2,
    width * scale,
    height * scale
  );
  const pixels = context.getImageData(0, 0, inputSize, inputSize).data;
  const tensorData = new Float32Array(3 * inputSize * inputSize);
  const planeSize = inputSize * inputSize;
  for (let index = 0; index < planeSize; index++) {
    tensorData[index] = pixels[index * 4] / 255;
    tensorData[planeSize + index] = pixels[index * 4 + 1] / 255;
    tensorData[planeSize * 2 + index] = pixels[index * 4 + 2] / 255;
  }

  const results = await model.run({ [inputName]: new Tensor("float32", tensorData, [1, 3, inputSize, inputSize]) });
  const output = results[outputName] as unknown as TensorLike;
  return extractBoxes(
    output,
    width,
    height,
    request.scoreThreshold ?? defaultScoreThreshold,
    request.iouThreshold ?? defaultIouThreshold
  );
}

parentPort?.on("message", async (request: DetectRequest) => {
  const started = Date.now();
  try {
    const boxes = await detect(request);
    const response: DetectResponse = { id: request.id, boxes, elapsedMs: Date.now() - started };
    parentPort?.postMessage(response);
  } catch (error) {
    const response: DetectResponse = {
      id: request.id,
      error: error instanceof Error ? error.message : String(error),
      elapsedMs: Date.now() - started
    };
    parentPort?.postMessage(response);
  }
});