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

declare global {
  interface Generator<T = unknown, TReturn = any, TNext = unknown> {
    select_many<TCollection>(
      selector: (item: T, index: number) => Generator<TCollection>
    ): Generator<TCollection, TReturn, TNext>;
    select_many<TCollection, TResult>(
      selector: (item: T, index: number) => Generator<TCollection>,
      result_selector: (item: T, part: TCollection) => TResult
    ): Generator<TResult, TReturn, TNext>;
  }

  interface AsyncGenerator<T = unknown, TReturn = any, TNext = unknown> {
    select_many<TCollection>(
      selector: (item: T, index: number) => AsyncGenerator<TCollection>
    ): AsyncGenerator<TCollection, TReturn, TNext>;
    select_many<TCollection, TResult>(
      selector: (item: T, index: number) => AsyncGenerator<TCollection>,
      result_selector: (item: T, part: TCollection) => Promised<TResult>
    ): AsyncGenerator<TResult, TReturn, TNext>;
  }
}

Build(
  "select_many",
  function* (selector, result_selector = (i, p) => p) {
    let index = 0;
    for (const item of this) {
      for (const p of selector(item, index).select((i) =>
        result_selector(item, i)
      )) {
        yield p;
      }
      index++;
    }
  },
  async function* (selector, result_selector = (i, p) => p) {
    let index = 0;
    for await (const item of this) {
      for await (const p of selector(item, index).select(
        async (i) => await Promise.resolve(result_selector(item, i))
      )) {
        yield p;
      }
      index++;
    }
  }
);