import type { Worker } from 'node:worker_threads';
import type { ChannelOptions } from './channel.ts';

declare const resultBrand: unique symbol;

/** An inert task descriptor. Carries the result type `T` and argument types `A`
 *  at the type level, but is not itself a `PromiseLike` or `AsyncIterable`.
 *  Dispatch via a {@link Runner} or `await` a live task returned by `mo()`. */
export type Task<T = unknown, A extends unknown[] = unknown[]> = {
  readonly uid: number;
  readonly id: string;
  readonly args: A;
  worker?: WorkerHandle;
  /** @internal Type brand for result type inference. Not present at runtime. */
  readonly [resultBrand]?: T;
};

/** Resolves the dispatch return type: `AsyncIterable<U>` for streaming tasks, `Promise<T>` otherwise. */
export type RunResult<T> = T extends AsyncIterable<infer U> ? AsyncIterable<U> : Promise<T>;

type TaskResults<T extends Task<any>[]> = {
  [K in keyof T]: T[K] extends Task<infer R> ? R : never;
};

/** A load balancing strategy for choosing which worker runs a task. */
export interface Balancer {
  /** Choose a worker for the given task. Called synchronously on every dispatch. */
  select(workers: readonly WorkerHandle[], task: Task): WorkerHandle;
  /** Optional cleanup on sync dispose. */
  [Symbol.dispose]?(): void;
  /** Optional cleanup on async dispose. */
  [Symbol.asyncDispose]?(): Promise<void>;
}

/** Options for configuring a worker pool. */
export interface WorkerOptions {
  /** Maximum time in ms to wait for in-flight tasks during async dispose. If exceeded, workers are force-terminated. */
  shutdownTimeout?: number;
  /** Load balancing strategy. Defaults to round-robin. */
  balance?: Balancer;
}

/** A handle to a specific worker in a pool. */
export interface WorkerHandle {
  /** Dispatches a task pinned to this worker. Returns `AsyncIterable<T>` for streaming tasks, `Promise<T>` otherwise. */
  exec<T>(task: Task<T>, opts?: ChannelOptions): RunResult<T>;
  /** The underlying worker thread. */
  readonly thread: Worker;
  /** Number of currently in-flight tasks on this worker. */
  readonly activeCount: number;
}

/**
 * A callable that dispatches tasks to a worker pool. Disposable via `using` or `[Symbol.dispose]()`.
 *
 * @param task - A single {@link Task} to run on a worker.
 * @returns `Promise<T>` for a single task, `Promise<[...results]>` for a batch, or `AsyncIterable<T>` for a streaming task.
 */
export type Runner = {
  /** Dispatches a batch of tasks in parallel and returns all results. */
  <T extends Task<any>[]>(tasks: [...T]): Promise<TaskResults<T>>;
  /** Dispatches a task. Returns `AsyncIterable<T>` for streaming tasks, `Promise<T>` otherwise. */
  <T>(task: Task<T>, opts?: ChannelOptions): RunResult<T>;
  /** AbortSignal that fires when the pool starts disposing. */
  readonly signal: AbortSignal;
  /** Read-only array of worker handles, one per pool worker. */
  readonly workers: readonly WorkerHandle[];
  /** Terminates all workers immediately. */
  [Symbol.dispose](): void;
  /** Aborts signal, waits for in-flight tasks to settle, then terminates workers. */
  [Symbol.asyncDispose](): Promise<void>;
};