import _ from "lodash";
import client, { Counter, Gauge, Histogram } from "prom-client";
import { ChoiceRequest } from "../models";

import logger from "../utils/logger";
import { DEFAULT_LABELS } from "../utils/RequestUtils";

export type EndFunction = (success: boolean, labels?: LabelValues) => void;
export type EndFunctionWithReq = (
  success: boolean,
  req: ChoiceRequest | null,
  labels?: LabelValues
) => void;

export class BucketSet {
  static readonly zeroTo100us: number[] = [1, 3, 5, 10, 30, 50, 100].map((b) => b / 1000000);
  static readonly zeroTo1ms: number[] = [0.01, 0.02, 0.03, 0.05, 0.1, 0.3, 0.5, 1].map(
    (b) => b / 1000
  );
  static readonly zeroTo5ms: number[] = [0.1, 0.3, 0.5, 0.7, 1, 1.5, 2, 3, 4, 5].map(
    (b) => b / 1000
  );
  static readonly zeroTo10ms: number[] = [0.1, 0.3, 0.5, 1, 1.5, 2, 3, 5, 7, 10].map(
    (b) => b / 1000
  );
  static readonly zeroTo25ms: number[] = [0.1, 0.5, 1, 2, 4, 7, 10, 12, 15, 20, 25].map(
    (b) => b / 1000
  );
  static readonly zeroTo100ms: number[] = [0.5, 1, 2, 3, 4, 5, 10, 25, 50, 100].map(
    (b) => b / 1000
  );
  static readonly zeroTo600ms: number[] = [1, 5, 10, 20, 50, 100, 150, 200, 300, 400, 600].map(
    (b) => b / 1000
  );
  static readonly zeroTo60s: number[] = [1, 4, 7, 10, 15, 20, 30, 45, 60];
  static readonly around100ms: number[] = [5, 20, 50, 80, 100, 130, 160, 200, 250, 300, 400].map(
    (b) => b / 1000
  );
}

export interface LabelValues {
  [key: string]: string | number | undefined;
}

class MetricService {
  private static instance: MetricService;
  private counters: Map<string, Counter<any>>;
  private histograms: Map<string, Histogram<any>>;
  private sections: Map<string, Section>;
  private gauges: Map<string, Gauge<any>>;

  constructor() {
    try {
      client.collectDefaultMetrics();
    } catch (err: any) {
      logger.warn(`[Prometheus] Failed to expose default nodejs metrics: ${err}`);
    }

    this.counters = new Map();
    this.histograms = new Map();
    this.sections = new Map();
    this.gauges = new Map();
  }

  public static getInstance(): MetricService {
    return this.instance || (this.instance = new MetricService());
  }

  /**
   * Adds a section. It automatically creates 2 counters (attempts and errors) and 1 histogram (section execution duration)
   * @param name
   * @param labelNames
   * @param buckets
   */
  public addSection(
    name: string,
    labelNames: string[] = [],
    buckets: number[] = BucketSet.zeroTo100ms
  ): void {
    if (!name) {
      throw new Error("Section name is required");
    }

    if (this.sections.has(name)) {
      return;
    }
    this.sections.set(name, new Section(name, addDefaultLabels(labelNames), buckets));
  }

  /**
   * Starts a section. I.e it increments the *attempt* counter, start the timer and returns a method that takes a
   * boolean (for success or failure)
   * @param name
   * @return {function(...[*]=)}
   */
  public startSection(name: string): EndFunction {
    const section = this.sections.get(name);
    return section!.start();
  }

  public startSectionWithReq(req: MaybeRequest, name: string): EndFunction {
    const section = this.sections.get(name);
    return section!.startWithRequest(req);
  }

  public startSectionWithDelayedReq(name: string): EndFunctionWithReq {
    const section = this.sections.get(name);
    return section!.startWithDelayedRequest();
  }

  public addCounter(name: string, help: string, labels?: string[]): void {
    if (this.counters.has(name)) {
      return;
    }
    this.counters.set(
      name,
      new client.Counter({
        name: name,
        help: help,
        labelNames: addDefaultLabels(labels),
      })
    );
  }

  public incCounter(name: string, number: number = 1, labels: LabelValues = {}): void {
    const counter = this.counters.get(name);
    counter!.inc(labels, number);
  }

  public incCounterWithReq(
    req: MaybeRequest,
    name: string,
    number: number = 1,
    labels: any = undefined
  ): void {
    labels = mergeDefaultLabelsFromRequest(req, labels);
    this.incCounter(name, number, labels);
  }

  public addGauge(
    name: string,
    help: string,
    labels?: string[],
    collect?: client.CollectFunction<Gauge<string>>
  ): void {
    if (this.gauges.has(name)) {
      return;
    }
    this.gauges.set(
      name,
      new client.Gauge({
        name: name,
        help: help,
        labelNames: addDefaultLabels(labels),
        collect,
      })
    );
  }

  public setGauge(name: string, number: number, labels: LabelValues = {}): void {
    const gauge = this.gauges.get(name);
    gauge!.labels(labels).set(number);
  }

  public setGaugeWithReq(
    req: MaybeRequest,
    name: string,
    number: number = 1,
    labels: any = undefined
  ): void {
    labels = mergeDefaultLabelsFromRequest(req, labels);
    this.setGauge(name, number, labels);
  }

  /**
   * Creates a Histogram with given name. The default bucket set is 0 to 100 (in ms)
   * @param {string} name
   * @param {string} help
   * @param labels
   * @param {number[]} buckets
   */
  public addHistogram(
    name: string,
    help: string,
    labels: string[] = [],
    buckets: number[] = BucketSet.zeroTo100ms
  ): void {
    if (this.histograms.has(name)) {
      return;
    }

    this.histograms.set(
      name,
      new client.Histogram({
        name: name,
        help: help,
        labelNames: addDefaultLabels(labels),
        // histogram buckets in ms
        buckets: buckets,
      })
    );
  }

  public observe(name: string, value: number, labels: LabelValues = {}): void {
    const histogram = this.histograms.get(name);
    histogram!.observe(labels, value);
  }

  public observeWithReq(
    req: MaybeRequest,
    name: string,
    value: number,
    labels: LabelValues = {}
  ): void {
    labels = mergeDefaultLabelsFromRequest(req, labels);
    this.observe(name, value, labels);
  }

  public startTimer(name: string): (labels?: LabelValues) => void {
    const histogram = this.histograms.get(name);
    return histogram!.startTimer();
  }

  public startTimerWithReq(req: MaybeRequest, name: string): (labels?: LabelValues) => void {
    const end = this.startTimer(name);

    return (labels?: LabelValues) => {
      labels = mergeDefaultLabelsFromRequest(req, labels);
      end(labels);
    };
  }

  public startTimerWithDelayedReq(name: string): (req: MaybeRequest, labels?: LabelValues) => void {
    const end = this.startTimer(name);

    return (req: MaybeRequest, labels?: LabelValues) => {
      labels = mergeDefaultLabelsFromRequest(req, labels);
      end(labels);
    };
  }
}

/**
 * Custom metric section.
 */
class Section {
  private histogram: Histogram<any>;

  constructor(
    name: string,
    labelNames: Omit<string, "error">[],
    bucketsInMs: number[] = BucketSet.zeroTo100ms
  ) {
    this.histogram = new client.Histogram({
      name: `${name}_duration_seconds`,
      help: `${name} section histogram`,
      labelNames: ["error", "error_code", ...(labelNames as string[])],
      // histogram buckets in ms
      buckets: bucketsInMs,
    });
  }

  start(): EndFunction {
    const end = this.histogram.startTimer();
    return (success, labels) => {
      labels = labels || {};
      labels.error = `${!success}`;

      end(labels);
    };
  }

  startWithRequest(req: MaybeRequest): EndFunction {
    const end = this.start();
    return (success, labels) => {
      labels = mergeDefaultLabelsFromRequest(req, labels);
      end(success, labels);
    };
  }

  startWithDelayedRequest(): EndFunctionWithReq {
    const end = this.start();
    return (success, req, labels) => {
      labels = mergeDefaultLabelsFromRequest(req, labels);
      end(success, labels);
    };
  }
}

export type MaybeRequest = Pick<ChoiceRequest, "allMetricsLabels"> | null;

const addDefaultLabels = (labels?: string[]): string[] => {
  return _.uniq([...(labels ?? []), ...DEFAULT_LABELS]);
};

const mergeDefaultLabelsFromRequest = (
  req: MaybeRequest,
  labels: LabelValues | undefined
): LabelValues => {
  labels = labels || {};
  return { ...labels, ...req?.allMetricsLabels };
};

export default MetricService;