import type { CryptoProvider } from './crypto.js';

/**
 * Signal Protocol KDF chain functions.
 *
 * These implement the key derivation logic used by both X3DH (initial root key)
 * and the Double Ratchet (root key ratchet + chain key ratchet).
 *
 * References:
 *  - Signal Double Ratchet spec, section 2.2 "KDF chains"
 *  - Signal X3DH spec, section 2.4 "Key derivation"
 */

// Info strings used in HKDF derivations (must match across all platforms)
const ROOT_KDF_INFO = new TextEncoder().encode('ShadeRootRatchet');
const CHAIN_KEY_CONSTANT = new Uint8Array([0x01]);
const MESSAGE_KEY_CONSTANT = new Uint8Array([0x02]);

/**
 * Root key ratchet step: given the current root key and a DH output,
 * derive a new root key and a new chain key.
 *
 * Uses HKDF with the DH output as IKM and the current root key as salt.
 * Output is 64 bytes: first 32 = new root key, last 32 = new chain key.
 */
export async function kdfRootKey(
  crypto: CryptoProvider,
  rootKey: Uint8Array,
  dhOutput: Uint8Array,
): Promise<{ newRootKey: Uint8Array; chainKey: Uint8Array }> {
  const derived = await crypto.hkdf(dhOutput, rootKey, ROOT_KDF_INFO, 64);
  return {
    newRootKey: derived.slice(0, 32),
    chainKey: derived.slice(32, 64),
  };
}

/**
 * Chain key ratchet step: derive the next chain key and a message key.
 *
 * Chain key → HMAC(chainKey, 0x01) = new chain key
 * Chain key → HMAC(chainKey, 0x02) = message key (used to encrypt one message)
 *
 * The message key is consumed (used once), the chain key advances.
 */
export async function kdfChainKey(
  crypto: CryptoProvider,
  chainKey: Uint8Array,
): Promise<{ newChainKey: Uint8Array; messageKey: Uint8Array }> {
  const [newChainKey, messageKey] = await Promise.all([
    crypto.hmacSha256(chainKey, CHAIN_KEY_CONSTANT),
    crypto.hmacSha256(chainKey, MESSAGE_KEY_CONSTANT),
  ]);
  return { newChainKey, messageKey };
}

/**
 * Derive the initial root key from X3DH shared secrets.
 *
 * Takes the concatenated DH outputs from X3DH (DH1 || DH2 || DH3 [|| DH4])
 * and derives a 32-byte root key using HKDF.
 *
 * Salt: 32 zero bytes (as per Signal spec)
 * Info: "ShadeX3DH"
 */
const X3DH_INFO = new TextEncoder().encode('ShadeX3DH');
const X3DH_SALT = new Uint8Array(32); // 32 zero bytes

export async function deriveInitialRootKey(
  crypto: CryptoProvider,
  sharedSecrets: Uint8Array[],
): Promise<Uint8Array> {
  // Concatenate all DH outputs
  const totalLength = sharedSecrets.reduce((sum, s) => sum + s.length, 0);
  const ikm = new Uint8Array(totalLength);
  let offset = 0;
  for (const secret of sharedSecrets) {
    ikm.set(secret, offset);
    offset += secret.length;
  }

  return crypto.hkdf(ikm, X3DH_SALT, X3DH_INFO, 32);
}