import { Worker } from "worker_threads";

function createWorker(data: Float32Array): Promise<Float32Array> {
  return new Promise((resolve, reject) => {
    const worker = new Worker("./vino.worker.js");
    worker.on("message", (result) => resolve(result as Float32Array));
    worker.on("error", reject);
    worker.on("exit", (code) => {
      if (code !== 0) {
        reject(new Error(`Worker stopped with exit code ${code}`));
      }
    });
    worker.postMessage(data);
  });
}

class TaskQueue {
  private concurrency: number;
  private queue: (() => void)[];
  private running: number;

  constructor(concurrency: number) {
    this.concurrency = concurrency;
    this.queue = [];
    this.running = 0;
  }

  private async runTask(task: () => Promise<any>): Promise<any> {
    this.running++;
    try {
      const result = await task();
      this.running--;
      this.next();
      return result;
    } catch (error) {
      this.running--;
      this.next();
      throw error;
    }
  }

  public addTask(task: () => Promise<any>): Promise<any> {
    return new Promise((resolve, reject) => {
      const run = () => this.runTask(task).then(resolve).catch(reject);
      this.queue.push(run);
      this.next();
    });
  }

  private next(): void {
    if (this.running < this.concurrency && this.queue.length > 0) {
      const task = this.queue.shift();
      task?.();
    }
  }
}

export async function cpuStartInfer() {
  const numTasks = 128;
  const patchSize = [1, 1, 80, 160, 160];
  const taskQueue = new TaskQueue(2); // 设置并发任务数，例如2

  console.time("总推理用时");

  try {
    const tasks = Array.from({ length: numTasks }, () => {
      const inferSample = new Float32Array(
        patchSize.reduce((a, b) => a * b)
      ).map(() => Math.random() * 2.0 - 1.0);
      return taskQueue.addTask(() => createWorker(inferSample));
    });

    const results = await Promise.all(tasks);
    console.log("所有任务完成，结果数量：", results.length);
  } catch (e) {
    console.error("推理失败:", e);
  }
  console.timeEnd("总推理用时");
}