import DataLoader from 'dataloader';

export type AutoHandledBatchFn<K, V> =
  (keys: ReadonlyArray<K>) => PromiseLike<ArrayLike<V | Error>> | ArrayLike<Promise<V | Error>>;

export type OrderFn<K, V> = (ids: ReadonlyArray<K>, values: (V | Error)[]) => (V | Error)[]

interface LoadEntityOptions<K, V, C = K> extends Omit<DataLoader.Options<K, V, C>, 'batch'> {
  id?: string;
  list?: boolean,
  loadKeyName?: keyof V;
  batch?: AutoHandledBatchFn<K, V>;
  fetch?: (key: K) => Promise<V | Error>;
  order?: OrderFn<K, V>;
  disablePrime?: boolean;
}

interface EntityLoaderOptions<V> {
  primeKey?: keyof V
}

export class EntityLoader<Entity, K = string, C = K> {
  private listLoaders: Map<string, DataLoader<K, Entity[]>> = new Map<string, DataLoader<K, Entity[]>>();
  private loaders: Map<string, DataLoader<K, Entity>> = new Map<string, DataLoader<K, Entity>>();

  constructor(private readonly options?: EntityLoaderOptions<Entity>) {}

  loaderName<Options extends LoadEntityOptions<K, Entity, C>>(options: Options) {
    return  options.loadKeyName ?
        options.loadKeyName as string :
        options.id ?
            options.id :
            this.constructor.name;
  }

  loadEntityList<Options extends LoadEntityOptions<K, Entity, C>>(options: Options) {

    const loaderName = this.loaderName(options);
    const cachedLoader = this.listLoaders.get(loaderName);

    if (cachedLoader) {
      return cachedLoader;
    }

    if (!options.batch && !options.fetch) {
      throw new Error('You must provide either batch or fetch function');
    }

    const loader = this.createLoader(options);

    this.loaders.set(loaderName, loader);

    return loader;
  }

  loadEntity<Options extends LoadEntityOptions<K, Entity, C>>(options: Options) {

    const loaderName = this.loaderName(options);
    const cachedLoader = this.loaders.get(loaderName);

    if (cachedLoader) {
      return cachedLoader;
    }

    if (!options.batch && !options.fetch) {
      throw new Error('You must provide either batch or fetch function');
    }

    const loader = this.createLoader(options);

    const _options = this.options;
    const _loaders = this.loaders;
    const _originalLoader = loader.load.bind(loader);

    const load: DataLoader<K, Entity, C>['load'] = async function(key: K) {
      const result = await _originalLoader(key);

      if (_options?.primeKey) {
        // fill cache for other loaders
        _loaders.forEach((loader, key) => {
          if (key !== loaderName && !options.disablePrime) {
            const primeKey = _options?.primeKey;
            if (result && primeKey && result[primeKey]) {
              const entityPrimeKey = result[primeKey];
              loader.prime(entityPrimeKey as any, result);
            }
          }
        });
      }
      return result;
    };

    loader.load = load;

    this.loaders.set(loaderName, loader);

    return loader;
  }

  protected createLoader<Options extends LoadEntityOptions<K, Entity, C>>(options: Options) {
    return new DataLoader<K, Entity, C>(async (ids) => {
      // handle batch
      if (options.batch) {
        return this.handleBatch<K, Entity, C>(options.batch, ids, options);
      }

      const fetchFn = options.fetch;

      if (!fetchFn) {
        throw new Error('You must provide either batch or fetch function');
      }

      return this.handleBatch<K, Entity, C>((ids) => {
        return ids.map(id => fetchFn(id));
      }, ids, options);

    }, {
      ...options,
      batch: true,
    });
  }

  protected async handleBatch<K, V, C>(
    batchFn: AutoHandledBatchFn<K, V>,
    ids: ReadonlyArray<K>,
    options?: Pick<LoadEntityOptions<K, V, C>, 'loadKeyName' | 'order'>
  ): Promise<ArrayLike<V | Error>> {

    const batchValues = batchFn(ids);
    let values: Array<V | Error> = [];

    if (Array.isArray(batchValues)) {
      // Make sure the error is returned
      values = await Promise.all<V | Error>(batchValues.map(batchValue => {
        return batchValue.catch((e: Error) => e);
      }));
    } else {
      values = await batchValues as Array<V | Error>;
    }

    if (options && options.order) {
      return options.order(ids, values);
    }

    if (options && options.loadKeyName) {
      return this.handleOrder(options.loadKeyName, ids, values);
    }

    return values;
  }

  protected handleOrder<K, V = Entity>(loadKeyName: keyof V, ids: ReadonlyArray<K>, values: (V | Error)[]) {
    return ids.map((id, index) => {
      const foundValue = values.find((value: (V | Error) & {[key: string]: any}) => {
        return (!(value instanceof Error) && value[loadKeyName] === id);
      });

      if (foundValue) {
        return foundValue;
      }

      if (values[index] instanceof Error) {
        return values[index];
      }

      return new Error(`Loader entity id ${id} not found`);
    });
  }
}