// Copyright Abridged, Inc. 2021,2024. All Rights Reserved.
// Node module: @collabland/aws
// This file is licensed under the MIT License.
// License text available at https://opensource.org/licenses/MIT

import {
  DataKeyPairSpec,
  DecryptRequest,
  EncryptRequest,
  GenerateDataKeyPairCommandOutput,
  GenerateDataKeyPairRequest,
  KMS,
  KMSClientConfig,
} from '@aws-sdk/client-kms';
import {
  ResourceType,
  SignOptions,
  getJwtIssuer,
  getJwtSecret,
  getResourceName,
  jwtSign,
  jwtVerify,
  loggers,
} from '@collabland/common';
import assert from 'assert';
import crypto from 'crypto';
import util from 'util';
import {getAWSKMSRegion, getDefaultConfig} from '../common/types.js';

const {trace, info} = loggers('collabland:aws:kms');

export type KMSServiceOptions = KMSClientConfig & {keyId?: string};

/**
 * KMS encrypted authentication token that extends JWT
 */
export type AEToken<T extends object = object> = T & {
  iat: number; // Issued at
  exp: number; // Expires at
  aud: string; // Audience
  iss: string; // Issuer
  sub: string; // Subject
};

export function getCMK() {
  return getResourceName(ResourceType.KMSKeyAlias);
}

export class KMSServiceClient {
  readonly keyId: string;
  readonly kms: KMS;

  constructor(options?: KMSServiceOptions) {
    const config = {
      ...getDefaultConfig(),
      region: getAWSKMSRegion(),
      ...options,
    };
    this.kms = new KMS(config);
    this.keyId = options?.keyId ?? getCMK();
    info('KMS region: %s, keyId: %s', config.region, this.keyId);
  }

  async encrypt(
    data: Uint8Array | string,
    ctx?: Record<string, string>,
  ): Promise<Buffer> {
    if (typeof data === 'string') {
      data = toUint8Array(data);
    }
    const req: EncryptRequest = {
      Plaintext: data,
      KeyId: this.keyId,
      EncryptionContext: ctx,
    };
    const res = await this.kms.encrypt(req);
    return Buffer.from(res.CiphertextBlob!);
  }

  async encryptJson(
    data: object,
    ctx?: Record<string, string>,
  ): Promise<Uint8Array> {
    return this.encrypt(JSON.stringify(data), ctx);
  }

  async decrypt(
    data: Uint8Array,
    ctx?: Record<string, string>,
  ): Promise<Buffer> {
    const req: DecryptRequest = {
      CiphertextBlob: data,
      KeyId: this.keyId,
      EncryptionContext: ctx,
    };
    const res = await this.kms.decrypt(req);
    return Buffer.from(res.Plaintext!);
  }

  async decryptJson<T extends object = object>(
    data: Uint8Array,
    ctx?: Record<string, string>,
  ): Promise<T> {
    const buffer = await this.decrypt(data, ctx);
    if (buffer == null) return buffer;
    const json = JSON.parse(fromUint8Array(buffer));
    trace('Decrypted json', json);
    return json;
  }

  async generateDataKeyPair(
    ctx?: Record<string, string>,
    keyPairSpec = DataKeyPairSpec.ECC_SECG_P256K1,
  ) {
    const req: GenerateDataKeyPairRequest = {
      KeyId: this.keyId,
      KeyPairSpec: keyPairSpec,
      EncryptionContext: ctx,
    };
    return this.kms.generateDataKeyPair(req);
  }

  /**
   * Export KMS key pair to PEM format
   * @param keyPair - KMS key pair
   * @returns
   */
  exportPEMKeys(keyPair: GenerateDataKeyPairCommandOutput) {
    const privateKeyObject = crypto.createPrivateKey({
      key: Buffer.from(keyPair.PrivateKeyPlaintext!),
      format: 'der',
      type: 'pkcs8',
    });

    const publicKeyObject = crypto.createPublicKey({
      key: Buffer.from(keyPair.PublicKey!),
      format: 'der',
      type: 'spki',
    });

    const privateKey = privateKeyObject.export({
      format: 'pem',
      type: 'pkcs8',
    });

    const publicKey = publicKeyObject.export({
      format: 'pem',
      type: 'spki',
    });
    return {privateKey, publicKey};
  }

  /**
   * Generate key pair in PEM format
   * @param ctx
   * @param keyPairSpec
   * @returns
   */
  async generatePEMKeyPair(
    ctx?: Record<string, string>,
    keyPairSpec = DataKeyPairSpec.ECC_SECG_P256K1,
  ) {
    const keyPair = await this.generateDataKeyPair(ctx, keyPairSpec);
    return this.exportPEMKeys(keyPair);
  }

  /**
   * Generate Mac-then-Encrypt authenticated encryption token (base64 encoded)
   * {@link https://en.wikipedia.org/wiki/Authenticated_encryption#MAC-then-Encrypt_(MtE) | MAC-then-Encrypt_(MtE)}
   *
   * 1. Load a jwt secret
   * 2. Encode/sign the data into a JWT
   * 3. Encrypt the JWT using KMS
   * 4. Encode the result in base 64 as the AE token
   *
   * @param data - JSON object
   * @param subject - JWT subject
   * @param expiresIn - JWT expiration
   */
  async generateAEToken<T extends object = object>(
    data: T,
    subject?: string,
    expiresIn?: number | string,
  ): Promise<string>;

  async generateAEToken<T extends object = object>(
    data: T,
    options?: SignOptions,
  ): Promise<string>;

  async generateAEToken<T extends object = object>(
    data: T,
    subject?: string | SignOptions,
    expiresIn?: number | string,
  ) {
    const signOptions: SignOptions =
      typeof subject === 'string' || subject == null ? {} : {};

    const secret = getJwtSecret();
    signOptions.issuer = signOptions.issuer ?? getJwtIssuer();

    if (typeof subject === 'string') {
      signOptions.subject = subject;
    }
    if (expiresIn != null) {
      signOptions.expiresIn = expiresIn;
    }

    const jwt = jwtSign(data, secret, signOptions);
    const buf = await this.encrypt(jwt);
    let str = buf!.toString('base64');
    // Make it base64 URL compatible. The `base64url` encoded string can still
    // be decoded using `Buffer.from(data, 'base64')`
    // '/' => '_', '+' => '-' and remove trailing '='
    str = str.replace(/\//g, '_').replace(/\+/g, '-').replace(/=/g, '');
    return str;
  }

  /**
   * Verify if the authenticated encryption (AE) token
   * {@link https://en.wikipedia.org/wiki/Authenticated_encryption#MAC-then-Encrypt_(MtE) | MAC-then-Encrypt_(MtE)}
   *
   * 1. Decode the token using base 64
   * 2. Decrypt the result using KMS
   * 3. Verify the JWT token
   *
   * @param token - Mac-then-Encrypt AE Token (base64 encoded)
   */
  async verifyAEToken<T extends object = object>(token: string) {
    trace('Verifying AE Token %s', token);
    const aeToken = Buffer.from(token, 'base64');
    const jwt = await this.decrypt(aeToken);
    trace('JWT: %s', jwt);
    assert(jwt != null);
    const secret = getJwtSecret();
    const data = jwtVerify(fromUint8Array(jwt), secret) as AEToken<T>;
    trace('AE Token', data);
    if (data.iss !== getJwtIssuer()) {
      throw new Error(`Invalid JWT issuer ${data.iss}`);
    }
    return data;
  }

  async generateJwtSecret(key = 'api.collab.land') {
    const buf = await this.encrypt(key);
    return buf!.toString('base64');
  }
}

export function toUint8Array(str?: string) {
  const encoder = new util.TextEncoder();
  return encoder.encode(str);
}

export function fromUint8Array(data?: Uint8Array) {
  const decoder = new util.TextDecoder('utf-8');
  return decoder.decode(data);
}