/**
 * Encryption Codec for Temporal Payload Encryption
 *
 * Implements AES-256-GCM encryption for workflow payloads.
 * Uses TEMPORAL_ENCODING_KEY environment variable.
 *
 * This is a copy of the worker's encryption codec - both need
 * to use the same encryption to communicate.
 */

import * as crypto from "crypto";
import {
  PayloadCodec,
  Payload,
  encodingKeys,
  METADATA_ENCODING_KEY,
} from "@temporalio/common";
import proto from "@temporalio/proto";

const { temporal } = proto;

const ENCODING_TYPE = "binary/encrypted";
const ALGORITHM = "aes-256-gcm";
const IV_LENGTH = 12;
const AUTH_TAG_LENGTH = 16;

/**
 * Derive a 256-bit key from the encoding key using SHA-256
 */
async function deriveKey(encodingKey: string): Promise<Buffer> {
  return crypto.createHash("sha256").update(encodingKey).digest();
}

/**
 * Encrypt data using AES-256-GCM
 */
function encrypt(data: Uint8Array, key: Buffer): Uint8Array {
  const iv = crypto.randomBytes(IV_LENGTH);
  const cipher = crypto.createCipheriv(ALGORITHM, key, iv);

  const encrypted = Buffer.concat([cipher.update(data), cipher.final()]);
  const authTag = cipher.getAuthTag();

  // Format: IV (12 bytes) + Ciphertext + AuthTag (16 bytes)
  return Buffer.concat([iv, encrypted, authTag]);
}

/**
 * Decrypt data using AES-256-GCM
 */
function decrypt(data: Uint8Array, key: Buffer): Uint8Array {
  const iv = data.slice(0, IV_LENGTH);
  const authTag = data.slice(-AUTH_TAG_LENGTH);
  const ciphertext = data.slice(IV_LENGTH, -AUTH_TAG_LENGTH);

  const decipher = crypto.createDecipheriv(ALGORITHM, key, iv);
  decipher.setAuthTag(authTag);

  return Buffer.concat([decipher.update(ciphertext), decipher.final()]);
}

/**
 * Encryption codec implementing Temporal's PayloadCodec interface
 */
export class EncryptionCodec implements PayloadCodec {
  private readonly key: Buffer;
  private readonly keyId: string;

  private constructor(key: Buffer, keyId: string) {
    this.key = key;
    this.keyId = keyId;
  }

  static async create(encodingKey: string): Promise<EncryptionCodec> {
    const key = await deriveKey(encodingKey);
    // Use first 8 chars of key hash as key ID
    const keyId = crypto
      .createHash("sha256")
      .update(encodingKey)
      .digest("hex")
      .slice(0, 8);
    return new EncryptionCodec(key, keyId);
  }

  async encode(payloads: Payload[]): Promise<Payload[]> {
    return payloads.map((payload) => {
      // Serialize the payload to bytes
      const bytes = temporal.api.common.v1.Payload.encode(payload).finish();

      // Encrypt
      const encrypted = encrypt(bytes, this.key);

      return {
        metadata: {
          [METADATA_ENCODING_KEY]: encodingKeys.METADATA_ENCODING_RAW,
          "encryption-codec": Buffer.from(ENCODING_TYPE),
          "encryption-key-id": Buffer.from(this.keyId),
        },
        data: encrypted,
      };
    });
  }

  async decode(payloads: Payload[]): Promise<Payload[]> {
    return payloads.map((payload) => {
      // Check if this payload is encrypted by us
      const encoding = payload.metadata?.["encryption-codec"];
      if (!encoding || Buffer.from(encoding).toString() !== ENCODING_TYPE) {
        return payload;
      }

      // Decrypt
      const decrypted = decrypt(payload.data!, this.key);

      // Deserialize back to Payload
      return temporal.api.common.v1.Payload.decode(decrypted);
    });
  }
}

/**
 * Create encryption codec from environment variable
 */
export async function createEncryptionCodec(): Promise<
  EncryptionCodec | undefined
> {
  const encodingKey = process.env.TEMPORAL_ENCODING_KEY;
  if (!encodingKey) {
    return undefined;
  }

  return EncryptionCodec.create(encodingKey);
}