// @flow strict
import { SYMBOL_ASYNC_ITERATOR } from '../polyfills/symbols';

import type { PromiseOrValue } from '../jsutils/PromiseOrValue';

/**
 * Given an AsyncIterable and a callback function, return an AsyncIterator
 * which produces values mapped via calling the callback function.
 */
export default function mapAsyncIterator<T, U>(
  iterable: AsyncIterable<T> | AsyncGenerator<T, void, void>,
  callback: (T) => PromiseOrValue<U>,
  rejectCallback?: (any) => PromiseOrValue<U>,
): AsyncGenerator<U, void, void> {
  // $FlowFixMe[prop-missing]
  const iteratorMethod = iterable[SYMBOL_ASYNC_ITERATOR];
  const iterator: any = iteratorMethod.call(iterable);
  let $return: any;
  let abruptClose;
  if (typeof iterator.return === 'function') {
    $return = iterator.return;
    abruptClose = (error: mixed) => {
      const rethrow = () => Promise.reject(error);
      return $return.call(iterator).then(rethrow, rethrow);
    };
  }

  function mapResult(result: IteratorResult<T, void>) {
    return result.done
      ? result
      : asyncMapValue(result.value, callback).then(iteratorResult, abruptClose);
  }

  let mapReject;
  if (rejectCallback) {
    // Capture rejectCallback to ensure it cannot be null.
    const reject = rejectCallback;
    mapReject = (error: mixed) =>
      asyncMapValue(error, reject).then(iteratorResult, abruptClose);
  }

  /* TODO: Flow doesn't support symbols as keys:
     https://github.com/facebook/flow/issues/3258 */
  return ({
    next(): Promise<IteratorResult<U, void>> {
      return iterator.next().then(mapResult, mapReject);
    },
    return() {
      return $return
        ? $return.call(iterator).then(mapResult, mapReject)
        : Promise.resolve({ value: undefined, done: true });
    },
    throw(error?: mixed): Promise<IteratorResult<U, void>> {
      if (typeof iterator.throw === 'function') {
        return iterator.throw(error).then(mapResult, mapReject);
      }
      return Promise.reject(error).catch(abruptClose);
    },
    [SYMBOL_ASYNC_ITERATOR]() {
      return this;
    },
  }: $FlowFixMe);
}

function asyncMapValue<T, U>(
  value: T,
  callback: (T) => PromiseOrValue<U>,
): Promise<U> {
  return new Promise((resolve) => resolve(callback(value)));
}

function iteratorResult<T>(value: T): IteratorResult<T, void> {
  return { value, done: false };
}