import { ServerlessCallback, ServerlessFunctionSignature, TwilioResponse } from "@twilio-labs/serverless-runtime-types/types";
import { PrometheusContentType, Registry, Pushgateway, Summary, Counter, register } from "prom-client";
import { DoormanLambdaContext } from "./DoormanHandlerContext";
import { shouldBlockRequest } from "../utils/blockUserAgent";
import { RequestOptions } from "https";
import { FastHttpPushgateway } from "../metrics/FastHttpPromGateway";
import '@twilio-labs/serverless-runtime-types';
import VoiceResponse from 'twilio/lib/twiml/VoiceResponse';

import { createLogger, format, Logger, transports } from "winston";
import LokiTransport from "winston-loki";
import pTimeout, { TimeoutError } from "promise.timeout";
import { ZodError } from "zod";
import { fromError } from "zod-validation-error";
import { setResponseJson } from "../utils/responseUtils";
import { BaseEvent, getSessionFromRequest, SessionType, setSessionOnResponse } from "./SessionHandler";

export type CallbackResult = Parameters<ServerlessCallback>;
export type FailFastCallback = () => CallbackResult;

export type DoormanLambda<T extends DoormanLambdaContext, U extends BaseEvent> = (
  context: Parameters<ServerlessFunctionSignature<T, U>>[0],
  event: Parameters<ServerlessFunctionSignature<T, U>>[1],
  callback: Parameters<ServerlessFunctionSignature<T, U>>[2],
  metricsRegistry: Registry<PrometheusContentType>,
  logger: Logger,
  // this is optional, but if called this method should provide an alternative callback provider that can fail fast
  failFastCallback: (fn: FailFastCallback) => void,
) => Promise<void>;

export enum CommonMetrics {
  RUNTIME = "FunctionRuntime",
  BLOCKED_REQUEST = "BlockedRequest",
  HTTP_CLIENT_ERROR = "HttpClientError",
  LOKI_ERROR = "LokiError",
  LOKI_LOG_AFTER_CLOSE = "LokiLogAfterClose",
  TWILIO_ISOLATION_BUSTED = "TwilioIsolationBusted",
  INNER_HANDLER_TIMEOUT = "InnerHandlerTimeout",
  INNER_HANDLER_MISSING_FAIL_FAST = "InnerHandlerMissingFailFast",
  RETURN_FALLBACK_RESPONSE = "ReturnFallbackResponse",
};

export function getMetricFromRegistry<T>(metricsRegistry: Registry, metric: string): T {
  return metricsRegistry.getSingleMetric(metric) as T;
};

export function gracefullyEndLogger(logger: Logger): Promise<void> {
  return new Promise((fulfill, reject) => {
    // Add logic for other transports here ...
    logger.on('finish', () => {
      fulfill();
    });

    logger.end();
  });
}

// TODO: make this response a voice message?
const REJECT_RESPONSE = new VoiceResponse();
REJECT_RESPONSE.reject();

const FALLBACK_CALLBACK: CallbackResult = [null, REJECT_RESPONSE];

const MINIMUM_MS_TO_SEND_RESPONSE: number = 250;
const FUNCTION_MAXIMUM_DURATION_MS: number = 10_000;
const INNER_HANDLER_MAXIMUM_DURATION_MS: number = 8_750;

/**
 * A decorator for twilio handlers. It provides a metrics registry and 
 * should implement timeout and cleanup jobs based on lambda timeout
 * @param handler 
 */
export function withMetrics<T extends DoormanLambdaContext, U extends BaseEvent>(
  functionName: string,
  handler: DoormanLambda<T, U>,
  parseSessionCookie?: boolean,
): ServerlessFunctionSignature<T, U> {
  return async (context, event, callback) => {
    console.log("[CommonHandler] created loki logger");
    console.log("[CommonHandler] clearing old metrics");
    register.clear();

    console.log("[CommonHandler] creating metrics registry");
    const metricsRegistry = new Registry();
    const requestOptions: RequestOptions = {
      headers: {
        'Authorization': 'Basic ' + Buffer.from(context.PUSHGATEWAY_USER + ':' + context.PUSHGATEWAY_PW).toString('base64'),
      },
    };

    const pushGateway = new FastHttpPushgateway(context.PUSHGATEWAY_URL, requestOptions, metricsRegistry);

    metricsRegistry.registerMetric(new Summary({
      name: CommonMetrics.RUNTIME,
      help: "Runtime of the function",
    }));

    metricsRegistry.registerMetric(new Counter({
      name: CommonMetrics.BLOCKED_REQUEST,
      help: "Blocked requests",
    }));

    metricsRegistry.registerMetric(new Counter({
      name: CommonMetrics.HTTP_CLIENT_ERROR,
      help: "Client side HTTP error",
      labelNames: ["ErrorCode"],
    }));

    metricsRegistry.registerMetric(new Counter({
      name: CommonMetrics.LOKI_ERROR,
      help: "Error connecting to loki for logs",
    }));

    metricsRegistry.registerMetric(new Counter({
      name: CommonMetrics.LOKI_LOG_AFTER_CLOSE,
      help: "Tried to write logs after log stream closed",
    }));

    metricsRegistry.registerMetric(new Counter({
      name: CommonMetrics.TWILIO_ISOLATION_BUSTED,
      help: "Twilio invocation lasted longer than 10s",
    }));

    metricsRegistry.registerMetric(new Counter({
      name: CommonMetrics.INNER_HANDLER_MISSING_FAIL_FAST,
      help: "Inner handler was timed out and we did not have any fallback response",
    }));

    metricsRegistry.registerMetric(new Counter({
      name: CommonMetrics.INNER_HANDLER_TIMEOUT,
      help: "Inner handler was timed out",
    }));

    metricsRegistry.registerMetric(new Counter({
      name: CommonMetrics.RETURN_FALLBACK_RESPONSE,
      help: "Common handler returned a fallback response",
    }));

    // Create a Winston logger
    const logger = createLogger({
      level: 'info',
      format: format.json(),
      transports: [
        new LokiTransport({
          host: context.LOKI_URL,
          labels: {
            service: "doorman",
            stage: context.STAGE,
            handler: functionName,
          },
          basicAuth: `${context.LOKI_USER}:${context.LOKI_PW}`,
          interval: 500,
          gracefulShutdown: true,
          clearOnError: true,
          batching: true,
          onConnectionError: (error) => {
            console.log("[CommonHandler] failed connecting to loki", error);
            getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.LOKI_ERROR).inc(1);
          },
        }),
        new transports.Console(), // Add other transports if needed
      ],
    });

    const orignalLog = {
      log: console.log,
      error: console.error,
      warn: console.warn,
      info: console.info,
    }

    // Override the base console log with winston
    console.log = function (...args) {
      if (logger.writable) {
        // @ts-ignore
        return logger.info.apply(logger, [...args]);
      }
      getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.LOKI_LOG_AFTER_CLOSE).inc(1);
      return orignalLog.log(...args);
    };
    console.error = function (...args) {
      if (logger.writable) {
        // @ts-ignore
        return logger.error.apply(logger, [...args]);
      }
      getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.LOKI_LOG_AFTER_CLOSE).inc(1);
      return orignalLog.error(...args);
    };
    console.warn = function (...args) {
      if (logger.writable) {
        // @ts-ignore
        return logger.warn.apply(logger, [...args]);
      }
      getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.LOKI_LOG_AFTER_CLOSE).inc(1);
      return orignalLog.warn(...args);
    };
    console.info = function (...args) {
      if (logger.writable) {
        // @ts-ignore
        return logger.info.apply(logger, [...args]);
      }
      getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.LOKI_LOG_AFTER_CLOSE).inc(1);
      return orignalLog.info(...args);
    };

    const summaryTimer = getMetricFromRegistry<Summary>(metricsRegistry, CommonMetrics.RUNTIME).startTimer();

    const startTime = Date.now();
    console.log(`[CommonHandler] started handler at ${startTime}`);

    if (parseSessionCookie) {
      console.log(`[CommonHandler] unwrapping session cookie`);
      event.request.session = await getSessionFromRequest(event, context);
    }

    const handlerResponsePromise: Promise<Parameters<ServerlessCallback>> = new Promise(async (resolve, reject) => {
      // intercept the callbackResult
      let callbackResult: Parameters<ServerlessCallback> | undefined;
      const tempCallback: ServerlessCallback = (err, payload) => {
        callbackResult = [err, payload];
      }

      // intercept the fail fast callback method, if it is provided
      let failFastCallbackMethod: FailFastCallback | undefined;
      const failFastCallbackProvider: (fn: FailFastCallback) => void = (fn) => {
        failFastCallbackMethod = fn;
      }

      // block requests before we even call the handler
      if (shouldBlockRequest(event)) {
        getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.BLOCKED_REQUEST).inc(1);
        const response = new Twilio.Response();
        response.setStatusCode(200);
        callbackResult = [null, response];
      } else {
        // wait for response for up to x ms, otherwise we call the failFastCallbackMethod
        try {
          await pTimeout(
            () => handler(context, event, tempCallback, metricsRegistry, logger, failFastCallbackProvider),
            INNER_HANDLER_MAXIMUM_DURATION_MS
          )();
        } catch (e) {
          console.log(e);
          if (e instanceof TimeoutError) {
            getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.INNER_HANDLER_TIMEOUT).inc(1);

            if (!failFastCallbackMethod) {
              getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.INNER_HANDLER_MISSING_FAIL_FAST).inc(1);
              reject("Timeout, but no failfast result was given");
              return;
            }

            callbackResult = failFastCallbackMethod();
          } else if (e instanceof ZodError) {
            // global catch for validation errors
            const response = new Twilio.Response();

            // return nice to read error message from ZOD
            setResponseJson(response, 400, {
              msg: fromError(e).toString(),
            });

            callbackResult = [null, response];
          }
        }
      }

      if (!callbackResult) {
        reject("No callback was given");
        return;
      }

      let statusCode: number | undefined = (callbackResult[1] as any)?.statusCode;

      if (statusCode && statusCode >= 400) {
        getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.HTTP_CLIENT_ERROR).inc({
          ErrorCode: statusCode,
        }, 1);
      }

      resolve(callbackResult);
    });

    console.time("[CommonHandler] nested handler time");

    const result = await handlerResponsePromise.catch((err) => {
      console.error("[CommonHandler] inner handler promise was rejected with reason: " + err);
      getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.RETURN_FALLBACK_RESPONSE).inc(1);
      return FALLBACK_CALLBACK;
    });

    console.timeEnd("[CommonHandler] nested handler time");

    const endTime = Date.now();
    const remainingTime = FUNCTION_MAXIMUM_DURATION_MS - (endTime - startTime);

    if (remainingTime < 0) {
      console.error("[CommonHandler] remaining time is negative for some reason...");
      getMetricFromRegistry<Counter>(metricsRegistry, CommonMetrics.TWILIO_ISOLATION_BUSTED).inc(1);
    }

    // abandoning metrics/logs, since there isn't enough time
    if (remainingTime <= MINIMUM_MS_TO_SEND_RESPONSE) {
      console.error("[CommonHandler] there is no time to send metrics / logs abandoning them");
      callback(...result);
      return;
    }

    console.log(`[CommonHandler] there is ${remainingTime} ms left to send metrics`);

    let metricsTimeout = setTimeout(() => {
      callback(...result);
    }, remainingTime - MINIMUM_MS_TO_SEND_RESPONSE);

    summaryTimer();
    console.log("[CommonHandler] attempting to push metrics...");

    try {
      // push logs
      await gracefullyEndLogger(logger);

      // push metrics
      await pushGateway.push({
        jobName: functionName,
        groupings: {
          stage: context.STAGE,
        },
      });
      console.log("[CommonHandler] pushed metrics / logs successfully");

    } catch (e: any) {
      console.log("[CommonHandler] failed to push metrics and or logs; quietly discarding them", e);
    }

    clearTimeout(metricsTimeout);
    callback(...result);
  };
};
