import { Build, Promised } from "./utils";

declare global {
  interface Generator<T = unknown, TReturn = any, TNext = unknown> {
    where(
      predicate: (item: T, index: number) => boolean
    ): Generator<T, TReturn, TNext>;
    where<TResult extends T>(
      predicate: (item: T, index: number) => item is TResult
    ): Generator<TResult, TReturn, TNext>;
  }

  interface AsyncGenerator<T = unknown, TReturn = any, TNext = unknown> {
    where(
      predicate: (item: T, index: number) => Promised<boolean>
    ): AsyncGenerator<T, TReturn, TNext>;
    where<TResult extends T>(
      predicate: (item: T, index: number) => item is TResult
    ): AsyncGenerator<TResult, TReturn, TNext>;
  }
}

Build(
  "where",
  function* (predicate) {
    let index = 0;
    for (const item of this) {
      if (predicate(item, index)) yield item;
      index++;
    }
  },
  async function* (predicate) {
    let index = 0;
    for await (const item of this) {
      if (await Promise.resolve(predicate(item, index))) yield item;
      index++;
    }
  }
);