import { Worker } from 'node:worker_threads';
import { setupWorker, execute, dispatchStream } from './execute.ts';

const workerEntryUrl = new URL(
  import.meta.url.endsWith('.ts') ? './worker-entry.ts' : './worker-entry.js',
  import.meta.url,
);
const workers = new Map<string, Worker>();
const activeCounts = new Map<Worker, number>();

function getWorker(id: string): Worker {
  let worker = workers.get(id);
  if (!worker) {
    worker = new Worker(workerEntryUrl);
    setupWorker(worker);
    worker.unref();
    workers.set(id, worker);
  }
  return worker;
}

function ref(worker: Worker): void {
  const count = activeCounts.get(worker) ?? 0;
  if (count === 0) worker.ref();
  activeCounts.set(worker, count + 1);
}

function unref(worker: Worker): void {
  const count = (activeCounts.get(worker) ?? 1) - 1;
  activeCounts.set(worker, count);
  if (count === 0) worker.unref();
}

export async function runOnDedicated<T>(id: string, args: unknown[]): Promise<T> {
  const worker = getWorker(id);
  ref(worker);
  try {
    return await execute<T>(worker, id, args);
  } catch (err) {
    if (err instanceof Error) {
      const Ctor = err.constructor as ErrorConstructor;
      throw new Ctor(err.message, { cause: err });
    }
    throw new Error(String(err), { cause: err });
  } finally {
    unref(worker);
  }
}

export function runStreamOnDedicated<T>(id: string, args: unknown[]): AsyncIterable<T> {
  const worker = getWorker(id);
  ref(worker);
  const { iterable, done } = dispatchStream<T>(worker, id, args);
  done.then(
    () => unref(worker),
    () => unref(worker),
  );
  return iterable;
}