import { Converter } from "../types.js";
import { Collector } from "./types.js";

export class SimpleCollector<TElement, TAccumulator, TResult> implements Collector<TElement, TAccumulator, TResult> {
  readonly #initialize: () => TAccumulator;
  readonly #accumulate: (accumulator: TAccumulator, element: TElement) => void;
  readonly #finalize: (accumulator: TAccumulator) => TResult;

  constructor(initialize: () => TAccumulator, accumulate: (accumulator: TAccumulator, element: TElement) => void,
    finalize: (accumulator: TAccumulator) => TResult) {
    this.#initialize = initialize;
    this.#accumulate = accumulate;
    this.#finalize = finalize;
  }

  initialize() {
    return this.#initialize();
  }

  accumulate(accumulator: TAccumulator, element: TElement) {
    this.#accumulate(accumulator, element);
  }

  finalize(accumulator: TAccumulator): TResult {
    return this.#finalize(accumulator);
  }
}

export class ToArrayCollector<TElement> implements Collector<TElement, TElement[], TElement[]> {
  initialize(): TElement[] {
    return [];
  }

  accumulate(accumulator: TElement[], element: TElement) {
    accumulator.push(element);
  }

  finalize(accumulator: TElement[]) {
    return accumulator;
  }
}

export class ToObjectCollector<TElement, TKey extends PropertyKey, TValue> implements Collector<TElement, Record<TKey, TValue>, Record<TKey, TValue>> {
  readonly #keySelector: Converter<TElement, TKey>;
  readonly #valueSelector: Converter<TElement, TValue>;

  constructor(keySelector: Converter<TElement, TKey>, valueSelector: Converter<TElement, TValue>) {
    this.#keySelector = keySelector;
    this.#valueSelector = valueSelector;
  }

  initialize() {
    return {} as Record<TKey, TValue>;
  }

  accumulate(accumulator: Record<TKey, TValue>, element: TElement) {
    const key = this.#keySelector(element);
    const value = this.#valueSelector(element);

    accumulator[key] = value;
  }

  finalize(accumulator: Record<TKey, TValue>) {
    return accumulator;
  }
}

export class ToObjectListCollector<TElement, TKey extends PropertyKey, TValue> implements Collector<TElement, Record<TKey, TValue[]>, Record<TKey, TValue[]>> {
  readonly #keySelector: Converter<TElement, TKey>;
  readonly #valueSelector: Converter<TElement, TValue>;

  constructor(keySelector: Converter<TElement, TKey>, valueSelector: Converter<TElement, TValue>) {
    this.#keySelector = keySelector;
    this.#valueSelector = valueSelector;
  }

  initialize() {
    return {} as Record<TKey, TValue[]>;
  }

  accumulate(accumulator: Record<TKey, TValue[]>, element: TElement) {
    const key = this.#keySelector(element);
    const value = this.#valueSelector(element);

    const l = accumulator[key];

    if (l) {
      l.push(value);
    } else {
      accumulator[key] = [value];
    }
  }

  finalize(accumulator: Record<TKey, TValue[]>) {
    return accumulator;
  }
}

export class ToMapCollector<TElement, TKey, TValue> implements Collector<TElement, Map<TKey, TValue>, Map<TKey, TValue>> {
  readonly #keySelector: Converter<TElement, TKey>;
  readonly #valueSelector: Converter<TElement, TValue>;

  constructor(keySelector: Converter<TElement, TKey>, valueSelector: Converter<TElement, TValue>) {
    this.#keySelector = keySelector;
    this.#valueSelector = valueSelector;
  }

  initialize() {
    return new Map<TKey, TValue>();
  }

  accumulate(accumulator: Map<TKey, TValue>, element: TElement) {
    const key = this.#keySelector(element);
    const value = this.#valueSelector(element);

    accumulator.set(key, value);
  }

  finalize(accumulator: Map<TKey, TValue>) {
    return accumulator;
  }
}

export class ToMapListCollector<TElement, TKey, TValue> implements Collector<TElement, Map<TKey, TValue[]>, Map<TKey, TValue[]>> {
  readonly #keySelector: Converter<TElement, TKey>;
  readonly #valueSelector: Converter<TElement, TValue>;

  constructor(keySelector: Converter<TElement, TKey>, valueSelector: Converter<TElement, TValue>) {
    this.#keySelector = keySelector;
    this.#valueSelector = valueSelector;
  }

  initialize() {
    return new Map<TKey, TValue[]>();
  }

  accumulate(accumulator: Map<TKey, TValue[]>, element: TElement) {
    const key = this.#keySelector(element);
    const value = this.#valueSelector(element);

    const l = accumulator.get(key);

    if (l) {
      l.push(value);
    } else {
      accumulator.set(key, [value]);
    }
  }

  finalize(accumulator: Map<TKey, TValue[]>) {
    return accumulator;
  }
}

export class ToSetCollector<TElement> implements Collector<TElement, Set<TElement>, Set<TElement>> {
  initialize() {
    return new Set<TElement>();
  }

  accumulate(accumulator: Set<TElement>, element: TElement) {
    accumulator.add(element);
  }

  finalize(accumulator: Set<TElement>) {
    return accumulator;
  }
}

export class JoinCollector implements Collector<any, any[], string> {
  readonly #delimiter: string;
  readonly #prefix: string;
  readonly #suffix: string;

  constructor(delimiter?: string, prefix?: string, suffix?: string) {
    this.#delimiter = delimiter ?? "";
    this.#prefix = prefix ?? "";
    this.#suffix = suffix ?? "";
  }

  initialize(): any[] {
    return [];
  }

  accumulate(accumulator: any[], element: string) {
    accumulator.push(element);
  }

  finalize(accumulator: any[]) {
    return this.#prefix + accumulator.join(this.#delimiter) + this.#suffix;
  }
}

export class SumCollector implements Collector<number, { sum: number; }, number> {
  initialize() {
    return { sum: 0 };
  }

  accumulate(accumulator: { sum: number; }, element: number) {
    accumulator.sum += element;
  }

  finalize(accumulator: { sum: number; }) {
    return accumulator.sum;
  }
}

export class BigIntSumCollector implements Collector<bigint, { sum: bigint; }, bigint> {
  initialize() {
    return { sum: 0n, };
  }

  accumulate(accumulator: { sum: bigint; }, element: bigint) {
    accumulator.sum += element;
  }

  finalize(accumulator: { sum: bigint; }) {
    return accumulator.sum;
  }
}

export class AverageCollector implements Collector<number, { count: number, sum: number; }, number> {
  initialize() {
    return { count: 0, sum: 0 };
  }

  accumulate(accumulator: { count: number, sum: number; }, element: number) {
    accumulator.count++;
    accumulator.sum += element;
  }

  finalize(accumulator: { count: number, sum: number; }) {
    return accumulator.count === 0 ? 0 : accumulator.sum / accumulator.count;
  }
}

export class BigIntAverageCollector implements Collector<bigint, { count: number, sum: bigint; }, bigint> {
  initialize() {
    return { count: 0, sum: 0n };
  }

  accumulate(accumulator: { count: number, sum: bigint; }, element: bigint) {
    accumulator.count++;
    accumulator.sum += element;
  }

  finalize(accumulator: { count: number, sum: bigint; }) {
    return accumulator.count === 0 ? 0n : accumulator.sum / BigInt(accumulator.count);
  }
}