import { registry, isModuleFrozen } from './registry.ts';
import { PromiseLikeTask } from './task.ts';
import { AsyncIterableTask } from './stream-task.ts';
import type { Task } from './runner.ts';

const counters = new Map<string, number>();

/**
 * Wraps a function to run on a worker thread. Must be called at module scope.
 * @param importMeta - The `import.meta` of the calling module, used to identify the source file.
 * @param fn - The function to offload to a worker thread.
 * @returns A callable that creates a {@link Task} when invoked.
 */

/** A value or a Task that resolves to that value on the worker. */
export type Arg<T> = T | Task<T>;

type TaskableArgs<A extends unknown[]> = { [K in keyof A]: Arg<A[K]> };

type Moroutine<A extends unknown[], R> = {
  (...args: A): Task<Awaited<R>, A> & PromiseLike<Awaited<R>>;
  (...args: TaskableArgs<A>): Task<Awaited<R>, A> & PromiseLike<Awaited<R>>;
  /** Stable id tying tasks back to this moroutine. Used by {@link isTask}. */
  readonly id: string;
};

type AsyncIterableMoroutine<A extends unknown[], Y> = {
  (...args: A): Task<AsyncIterable<Y>, A> & AsyncIterable<Y>;
  (...args: TaskableArgs<A>): Task<AsyncIterable<Y>, A> & AsyncIterable<Y>;
  /** Stable id tying tasks back to this moroutine. Used by {@link isTask}. */
  readonly id: string;
};

type IsNever<T> = [T] extends [never] ? true : false;
type MoReturn<A extends unknown[], R> =
  IsNever<R> extends true
    ? Moroutine<A, R>
    : [R] extends [AsyncGenerator<infer Y, any, any>]
      ? AsyncIterableMoroutine<A, Y>
      : Moroutine<A, R>;

function isAsyncGeneratorFunction(fn: Function): boolean {
  return fn.constructor.name === 'AsyncGeneratorFunction';
}

export function mo<A extends unknown[], R>(importMeta: ImportMeta, fn: (...args: A) => R): MoReturn<A, R> {
  const url = importMeta.url;

  if (isModuleFrozen(url)) {
    throw new Error(
      `Cannot call mo() for ${url} after a task from this module has been dispatched. ` +
        'All mo() calls must happen at module load time.',
    );
  }

  const index = counters.get(url) ?? 0;
  counters.set(url, index + 1);
  const id = `${url}#${index}`;

  registry.set(id, fn);

  const factory = isAsyncGeneratorFunction(fn)
    ? (...args: unknown[]) => new AsyncIterableTask(id, args)
    : (...args: unknown[]) => new PromiseLikeTask<R>(id, args);
  Object.defineProperty(factory, 'id', { value: id, enumerable: true });
  return factory as unknown as MoReturn<A, R>;
}