import AWS from "aws-sdk";
import https from "https";
import ShutdownService from "../../services/shutdown.service";
import MetricService from "../../services/metric.service";
import logger from "../../utils/logger";
import { config } from "../../utils/config";
import constants from "../../utils/constants";
import { createAZAwareDAXClient } from "@ogury/dax-client-az-aware";
import { OgyCache, OgyCacheBatchValueGetterResult } from "@ogury/ogy-cache";

const CACHE_REQUESTS_METRIC = "dynamo_cache_requests_total";
const CACHE_SIZE_METRIC = "dynamo_cache_size_total";

const getDynamoConfig = (): AWS.DynamoDB.Types.ClientConfiguration => {
  const isMockEnv = config.MOCK_ENVS.split(",")
    .map((val) => val.trim())
    .includes(config.OGURY_ENV);

  const configuration: AWS.DynamoDB.Types.ClientConfiguration = {
    region: config.AWS_REGION,
  };

  if (isMockEnv) {
    configuration.endpoint = config.LOCAL_DYNAMO_ENDPOINT;
  } else {
    configuration.httpOptions = {
      agent: new https.Agent({ rejectUnauthorized: true, keepAlive: true }),
    };
  }

  return configuration;
};

export class DDBConnector {
  client: Promise<AWS.DynamoDB>;
  cache: OgyCache<string, any>;

  constructor(private readonly tableName: string, daxEndpoint: string | undefined = undefined) {
    const daxEnabled = !!daxEndpoint;

    const dynamoConfig = getDynamoConfig();

    if (daxEnabled) {
      const daxConfig = {
        ...dynamoConfig,
        endpoint: daxEndpoint,
      };

      this.client = createAZAwareDAXClient(daxConfig, false) as unknown as Promise<AWS.DynamoDB>;

      logger.info(`DAX client configured at ${daxEndpoint} for table ${tableName}`);
    } else {
      this.client = Promise.resolve(new AWS.DynamoDB(dynamoConfig));

      logger.info(`Dynamo client configured for table ${tableName}`);
    }

    this.cache = new OgyCache(this.getItemsFromDynamo.bind(this), {
      maxItemsPerBatch: 50,
      getterTimeout: 15000,
      staleAfter: { value: 9 * 60 * 1000, jitter: 60 * 1000 },
      expireAfter: { value: 14 * 60 * 1000, jitter: 60 * 1000 },
      cacheNulls: false,
      logError: logger.error,
    });

    ShutdownService.addShutdownListener(async () => {
      if (daxEnabled) {
        logger.info(`Closing DAX ${tableName}`);
        const client = await this.getClient();
        (client as any).shutdown();
      }
    });

    this.setupCacheMonitoring();
  }

  setupCacheMonitoring() {
    const metricService = MetricService.getInstance();
    const tableName = this.tableName;
    const cache = this.cache;

    metricService.addGauge(
      CACHE_REQUESTS_METRIC,
      "DynamoDB cache hits, misses and stales",
      ["table_name", "response_type"],
      function () {
        const stats = cache.stats;
        this.set(
          {
            table_name: tableName,
            response_type: "hit",
          },
          stats.itemRequests.hitsFresh
        );
        this.set(
          {
            table_name: tableName,
            response_type: "stale",
          },
          stats.itemRequests.hitsStale
        );
        this.set(
          {
            table_name: tableName,
            response_type: "miss",
          },
          stats.itemRequests.misses
        );
      }
    );

    metricService.addGauge(CACHE_SIZE_METRIC, "DynamoDB cache size", ["table_name"], function () {
      this.set(
        {
          table_name: tableName,
        },
        cache.stats.cacheSize
      );
    });
  }

  async getClient(): Promise<AWS.DynamoDB> {
    return this.client;
  }

  async get<T>(key: string): Promise<T | null> {
    return this.cache.get(key);
  }

  async batchGet<T>(keys: string[]): Promise<Record<string, T | null>> {
    const items = await this.cache.batchGet(keys);
    const result: Record<string, T | null> = {};
    for (let i = 0; i < keys.length; ++i) {
      result[keys[i]] = items[i];
    }
    return result;
  }

  protected async getItemsFromDynamo(
    keys: string[]
  ): Promise<OgyCacheBatchValueGetterResult<string, any>> {
    return this.internalBatchGet(keys);
  }

  protected async internalBatchGet<T>(
    keys: string[]
  ): Promise<OgyCacheBatchValueGetterResult<string, any>> {
    const items: OgyCacheBatchValueGetterResult<string, any> = [];

    try {
      const query = {
        RequestItems: {
          [this.tableName]: {
            Keys: keys.map((k) => ({
              _id: { S: k },
            })),
          },
        },
      };

      const client = await this.getClient();
      const result = await client.batchGetItem(query).promise();

      for (const r of result.Responses?.[this.tableName] ?? []) {
        const key = r._id.S!;
        const item = this.processItem<T>(key, r);

        if (item == null) {
          throw new Error(`Invalid item: ${key}`);
        }

        items.push([key, item]);
      }

      return items;
    } catch (err: any) {
      logger.error(err, `Error getting keys: ${keys}`);
      err.key = keys;
      throw err;
    }
  }

  processItem<T>(key: string, item: any): T | null {
    if (item == null) {
      if (key !== constants.HEALTH_KEY) {
        logger.debug(() => `Key not found: ${key}`);
      }

      return null;
    }

    if (item._ttl) {
      const currentTimeInMs = Math.round(Date.now() / 1000);
      const itemTTL = parseInt(item._ttl.N);

      if (currentTimeInMs > itemTTL) {
        return null;
      }
    }

    if (item.counter) {
      return parseInt(item.counter.N) as unknown as T;
    }

    return JSON.parse(item.value.S);
  }

  async save<T>(key: string, data: T, options = { expiry: 86400 }): Promise<T> {
    const item: AWS.DynamoDB.Types.PutItemInput = {
      Item: {
        _id: { S: key },
        value: { S: JSON.stringify(data) },
      },
      TableName: this.tableName,
    };

    if (options.expiry) {
      item.Item._ttl = {
        N: `${Math.round(Date.now() / 1000) + options.expiry}`,
      };
    }

    try {
      const client = await this.getClient();
      await client.putItem(item).promise();
      return data;
    } catch (err: any) {
      logger.error(err, `Error saving key ${key}`);
      throw err;
    }
  }

  async counter(key: string, delta: number, expiry?: number | null): Promise<number> {
    const item: AWS.DynamoDB.Types.UpdateItemInput = {
      Key: { _id: { S: key } },
      ReturnValues: "ALL_NEW",
      ExpressionAttributeNames: {
        "#counter": "counter",
      },
      ExpressionAttributeValues: {
        ":incr": { N: `${delta}` },
        ":initial_value": { N: `${1 - delta}` },
      },
      UpdateExpression:
        expiry != null
          ? "SET #counter = if_not_exists (#counter, :initial_value) + :incr, #ttl = :ttl"
          : "SET #counter = if_not_exists (#counter, :initial_value) + :incr",
      TableName: this.tableName,
    };
    if (expiry != null) {
      const ttl = Math.round(Date.now() / 1000) + expiry;
      item.ExpressionAttributeNames!["#ttl"] = "_ttl";
      item.ExpressionAttributeValues![":ttl"] = { N: `${ttl}` };
    }

    const client = await this.getClient();
    const res = await client.updateItem(item).promise();
    return parseFloat(res.Attributes?.counter.N ?? "");
  }

  async health(): Promise<boolean> {
    try {
      await this.internalBatchGet([constants.HEALTH_KEY]);
      return true;
    } catch (err: any) {
      return false;
    }
  }
}