/**
 * Temporal Payload Encryption Codec
 *
 * Encrypts workflow payloads using AES-256-GCM following Temporal's
 * recommended pattern: serialize the entire payload with Protobuf,
 * then encrypt.
 *
 * Reference: https://github.com/temporalio/samples-typescript/blob/main/encryption/src/encryption-codec.ts
 */

import { webcrypto as crypto } from "node:crypto";
import {
  METADATA_ENCODING_KEY,
  Payload,
  PayloadCodec,
  ValueError,
} from "@temporalio/common";
import proto from "@temporalio/proto";
import { logger } from "@mesh-tech/logger";

const { temporal } = proto;

const ENCODING = "binary/encrypted";
const METADATA_ENCRYPTION_KEY_ID = "encryption-key-id";
const IV_LENGTH = 12;

/**
 * Derive a CryptoKey from an encoding key string using SHA-256
 */
async function deriveKey(encodingKey: string): Promise<crypto.CryptoKey> {
  const keyData = await crypto.subtle.digest(
    "SHA-256",
    new TextEncoder().encode(encodingKey)
  );
  return crypto.subtle.importKey("raw", keyData, { name: "AES-GCM" }, false, [
    "encrypt",
    "decrypt",
  ]);
}

/**
 * Encrypt data using AES-256-GCM
 * Output: IV (12 bytes) + Ciphertext+AuthTag (WebCrypto format)
 */
async function encrypt(
  data: Uint8Array,
  key: crypto.CryptoKey
): Promise<Uint8Array> {
  const iv = crypto.getRandomValues(new Uint8Array(IV_LENGTH));
  const encrypted = await crypto.subtle.encrypt(
    { name: "AES-GCM", iv },
    key,
    data
  );
  // Prepend IV to ciphertext+authTag
  const result = new Uint8Array(IV_LENGTH + encrypted.byteLength);
  result.set(iv, 0);
  result.set(new Uint8Array(encrypted), IV_LENGTH);
  return result;
}

/**
 * Decrypt data using AES-256-GCM
 * Input: IV (12 bytes) + Ciphertext+AuthTag
 */
async function decrypt(
  data: Uint8Array,
  key: crypto.CryptoKey
): Promise<Uint8Array> {
  const iv = data.slice(0, IV_LENGTH);
  const ciphertextWithTag = data.slice(IV_LENGTH);
  const decrypted = await crypto.subtle.decrypt(
    { name: "AES-GCM", iv },
    key,
    ciphertextWithTag
  );
  return new Uint8Array(decrypted);
}

/**
 * Payload codec that encrypts/decrypts payloads using AES-256-GCM.
 *
 * Uses Protobuf serialization for the entire payload (data + metadata)
 * following Temporal's recommended pattern.
 */
export class EncryptionCodec implements PayloadCodec {
  private constructor(
    private readonly key: crypto.CryptoKey,
    private readonly keyId: string
  ) {}

  static async create(encodingKey: string): Promise<EncryptionCodec> {
    const key = await deriveKey(encodingKey);
    // Use truncated hash as key ID for identification
    const keyIdData = await crypto.subtle.digest(
      "SHA-256",
      new TextEncoder().encode(encodingKey)
    );
    const keyId = Buffer.from(keyIdData).toString("hex").slice(0, 16);
    return new EncryptionCodec(key, keyId);
  }

  async encode(payloads: Payload[]): Promise<Payload[]> {
    return Promise.all(
      payloads.map(async (payload) => ({
        metadata: {
          [METADATA_ENCODING_KEY]: new TextEncoder().encode(ENCODING),
          [METADATA_ENCRYPTION_KEY_ID]: new TextEncoder().encode(this.keyId),
        },
        // Serialize entire payload with Protobuf, then encrypt
        data: await encrypt(
          temporal.api.common.v1.Payload.encode(payload).finish(),
          this.key
        ),
      }))
    );
  }

  async decode(payloads: Payload[]): Promise<Payload[]> {
    return Promise.all(
      payloads.map(async (payload) => {
        // Check if encrypted
        const encoding = payload.metadata?.[METADATA_ENCODING_KEY];
        if (!encoding) {
          return payload;
        }

        const encodingStr = new TextDecoder().decode(encoding);
        if (encodingStr !== ENCODING) {
          return payload;
        }

        if (!payload.data) {
          throw new ValueError("Payload data is missing");
        }

        // Decrypt and deserialize with Protobuf
        const decryptedBytes = await decrypt(payload.data, this.key);
        return temporal.api.common.v1.Payload.decode(decryptedBytes);
      })
    );
  }
}

/**
 * Creates an encryption codec from TEMPORAL_ENCODING_KEY env var.
 * Returns undefined if not set.
 */
export async function createEncryptionCodec(): Promise<EncryptionCodec | undefined> {
  const encodingKey = process.env.TEMPORAL_ENCODING_KEY;
  if (!encodingKey) {
    logger.warn("No TEMPORAL_ENCODING_KEY set, payloads will not be encrypted");
    return undefined;
  }
  logger.info("Payload encryption enabled");
  return EncryptionCodec.create(encodingKey);
}