/**
 * Bounded-concurrency async map.
 *
 * Pulls items from an `AsyncIterable` source (lazily — never materializing
 * the whole sequence) and runs `fn` on at most `concurrency` of them at
 * once. Failures bubble unless `continueOnError` is set, in which case they
 * are reported via `onError` and the pool keeps draining.
 */
import { CancelledError } from '../schemas/errors.js';

export interface ConcurrentMapOptions<T> {
  concurrency: number;
  signal?: AbortSignal;
  continueOnError?: boolean;
  onError?: (item: T, err: unknown) => void;
}

export async function runWithConcurrency<T>(
  source: AsyncIterable<T>,
  fn: (item: T) => Promise<void>,
  opts: ConcurrentMapOptions<T>,
): Promise<void> {
  if (opts.concurrency < 1) throw new Error('concurrency must be ≥ 1');

  const iter = source[Symbol.asyncIterator]();
  const inFlight = new Set<Promise<void>>();
  let firstError: unknown = null;
  let aborted = false;

  function checkAbort(): void {
    if (opts.signal?.aborted) {
      aborted = true;
      const reason = opts.signal.reason;
      throw reason instanceof Error ? reason : new CancelledError(String(reason ?? 'aborted'));
    }
  }

  async function pumpOne(): Promise<boolean> {
    checkAbort();
    const next = await iter.next();
    if (next.done === true) return false;
    const item = next.value;
    const task = (async () => {
      try {
        await fn(item);
      } catch (err) {
        if (opts.continueOnError === true) {
          opts.onError?.(item, err);
        } else if (firstError === null) {
          firstError = err;
        }
      }
    })().finally(() => {
      inFlight.delete(task);
    });
    inFlight.add(task);
    return true;
  }

  try {
    // Prime the pool
    while (inFlight.size < opts.concurrency) {
      if (firstError !== null || aborted) break;
      const more = await pumpOne();
      if (!more) break;
    }
    // Maintain saturation
    while (inFlight.size > 0) {
      if (firstError !== null && opts.continueOnError !== true) break;
      await Promise.race(inFlight);
      while (
        inFlight.size < opts.concurrency &&
        firstError === null &&
        !aborted
      ) {
        const more = await pumpOne();
        if (!more) break;
      }
    }
    // Drain whatever survived
    if (inFlight.size > 0) await Promise.allSettled(inFlight);
  } finally {
    if (typeof iter.return === 'function') {
      await iter.return().catch(() => undefined);
    }
  }

  if (firstError !== null && opts.continueOnError !== true) {
    throw firstError;
  }
}
