import { createRequire } from "node:module";
import os from "node:os";
import path from "node:path";
import type { IdDrawStyle } from "./id-config";

export type WorkerFont = {
  path: string;
  family: string;
};

export type GenerateTask = {
  text: string;
  outPath: string;
};

const nodeRequire = createRequire(path.join(process.cwd(), "package.json"));

function spawnWorker(workerPath: string, workerData: unknown) {
  const threads = nodeRequire("node:worker_threads") as typeof import("node:worker_threads");
  return new threads.Worker(workerPath, { workerData });
}

function workerCount(taskCount: number): number {
  const cpus =
    typeof os.availableParallelism === "function"
      ? os.availableParallelism()
      : os.cpus().length;
  return Math.max(1, Math.min(cpus, 8, taskCount));
}

function splitTasks<T>(items: T[], parts: number): T[][] {
  const chunks: T[][] = Array.from({ length: parts }, () => []);
  items.forEach((item, index) => {
    chunks[index % parts]?.push(item);
  });
  return chunks.filter((chunk) => chunk.length > 0);
}

export async function generatePngsInWorkers(options: {
  imagePath: string;
  fonts: WorkerFont[];
  style: IdDrawStyle;
  tasks: GenerateTask[];
  onProgress: (completed: number, total: number) => void;
}): Promise<void> {
  const { imagePath, fonts, style, tasks, onProgress } = options;
  if (tasks.length === 0) return;

  const workerPath = path.join(process.cwd(), "workers", "generate-id.cjs");
  const chunks = splitTasks(tasks, workerCount(tasks.length));
  let completed = 0;

  await Promise.all(
    chunks.map(
      (chunk) =>
        new Promise<void>((resolve, reject) => {
          const worker = spawnWorker(workerPath, {
            imagePath,
            fonts,
            style,
            tasks: chunk,
          });

          worker.on("message", (message: { type?: string }) => {
            if (message?.type !== "progress") return;
            completed += 1;
            onProgress(completed, tasks.length);
          });

          worker.once("error", reject);
          worker.once("exit", (code) => {
            if (code === 0) {
              resolve();
              return;
            }
            reject(new Error(`Image worker stopped with code ${code}`));
          });
        }),
    ),
  );
}