import { Converter } from "./sync.js";

export interface Collector<TElement, TAccumulator, TResult> {
  initialize(): TAccumulator;
  accumulate(accumulator: TAccumulator, element: TElement): void;
  finalize(accumulator: TAccumulator): TResult;
}

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 function create<TElement, TAccumulator, TResult>(initialize: () => TAccumulator, accumulate: (accumulator: TAccumulator, element: TElement) => void,
  finalize: (accumulator: TAccumulator) => TResult): Collector<TElement, TAccumulator, TResult> {
  return new SimpleCollector(initialize, accumulate, finalize);
}

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

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

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

const toArrayCollector = new ToArrayCollector<any>();

export function toArray<TElement>(): Collector<TElement, any, TElement[]> {
  return toArrayCollector;
}

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 function toObject<TElement, TKey extends PropertyKey, TValue>(keySelector: Converter<TElement, TKey>, valueSelector: Converter<TElement, TValue>): Collector<TElement, any, Record<TKey, TValue>> {
  return new ToObjectCollector<TElement, TKey, TValue>(keySelector, valueSelector);
}

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 function toMap<TElement, TKey, TValue>(keySelector: Converter<TElement, TKey>, valueSelector: Converter<TElement, TValue>): Collector<TElement, any, Map<TKey, TValue>> {
  return new ToMapCollector<TElement, TKey, TValue>(keySelector, valueSelector);
}

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;
  }
}

const toSetCollector = new ToSetCollector<any>();

export function toSet<TElement>(): Collector<TElement, any, Set<TElement>> {
  return toSetCollector;
}

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 function join(delimiter?: string, prefix?: string, suffix?: string): Collector<any, any, string> {
  return new JoinCollector(delimiter, prefix, suffix);
}

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;
  }
}

const sumCollector = new SumCollector();

export function sum(): Collector<number, any, number> {
  return sumCollector;
}

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;
  }
}

const bigintSumCollector = new BigIntSumCollector();

export function bigintSum(): Collector<bigint, any, bigint> {
  return bigintSumCollector;
}

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;
  }
}

const averageCollector = new AverageCollector();

export function average(): Collector<number, any, number> {
  return averageCollector;
}

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);
  }
}

const bigintAverageCollector = new BigIntAverageCollector();

export function bigintAverage(): Collector<bigint, any, bigint> {
  return bigintAverageCollector;
}