import { markTaskBatch } from '../runtime/awaitedTasks';

let batchId = 0;

const getBatchId = () => batchId++;

export const opMarker = Symbol('opMarker');

const wrapPromiseUtil =
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  <T extends (values: any) => any>(nativePromiseUtil: T): T =>
    ((values: Parameters<T>) => {
      let argsWithMarker = 0;
      values.forEach((arg) => {
        if (arg instanceof Promise && (arg as Promise<unknown> & { [opMarker]: boolean })[opMarker]) {
          argsWithMarker++;
        }
      });

      if (argsWithMarker > 0) {
        markTaskBatch(argsWithMarker, getBatchId());
      }

      // eslint-disable-next-line @typescript-eslint/no-unsafe-return
      return nativePromiseUtil(values);
    }) as T;

const nativePromiseAll = Promise.all.bind(Promise);
const nativePromiseRace = Promise.race.bind(Promise);
const nativePromiseAny = Promise.any.bind(Promise);
const nativePromiseSettled = Promise.allSettled.bind(Promise);
const nativePromiseAllSettled = Promise.allSettled.bind(Promise);

globalThis.Promise.all = wrapPromiseUtil(nativePromiseAll);
globalThis.Promise.race = wrapPromiseUtil(nativePromiseRace);
globalThis.Promise.any = wrapPromiseUtil(nativePromiseAny);
globalThis.Promise.allSettled = wrapPromiseUtil(nativePromiseSettled);
globalThis.Promise.allSettled = wrapPromiseUtil(nativePromiseAllSettled);
