package no.zyon.shade.protocol

import no.zyon.shade.crypto.CryptoProvider

/**
 * KDF chain functions for the Signal Protocol ratchet.
 *
 * MUST produce byte-identical output to @shade/core/keys.ts.
 * Info strings and salts are fixed constants and must not change.
 */

// Must match the TypeScript version EXACTLY
private val ROOT_KDF_INFO = "ShadeRootRatchet".toByteArray(Charsets.UTF_8)
private val CHAIN_KEY_CONSTANT = byteArrayOf(0x01)
private val MESSAGE_KEY_CONSTANT = byteArrayOf(0x02)

private val X3DH_INFO = "ShadeX3DH".toByteArray(Charsets.UTF_8)
private val X3DH_SALT = ByteArray(32) // 32 zero bytes

data class RootKdfResult(val newRootKey: ByteArray, val chainKey: ByteArray)
data class ChainKdfResult(val newChainKey: ByteArray, val messageKey: ByteArray)

/**
 * Root key ratchet step.
 * HKDF(ikm=dhOutput, salt=rootKey, info="ShadeRootRatchet", length=64)
 * → first 32 bytes = new root key, last 32 bytes = chain key
 */
fun kdfRootKey(crypto: CryptoProvider, rootKey: ByteArray, dhOutput: ByteArray): RootKdfResult {
    val derived = crypto.hkdf(dhOutput, rootKey, ROOT_KDF_INFO, 64)
    return RootKdfResult(
        newRootKey = derived.copyOfRange(0, 32),
        chainKey = derived.copyOfRange(32, 64),
    )
}

/**
 * Chain key ratchet step.
 * HMAC(chainKey, 0x01) = new chain key
 * HMAC(chainKey, 0x02) = message key (used once)
 */
fun kdfChainKey(crypto: CryptoProvider, chainKey: ByteArray): ChainKdfResult {
    val newChainKey = crypto.hmacSha256(chainKey, CHAIN_KEY_CONSTANT)
    val messageKey = crypto.hmacSha256(chainKey, MESSAGE_KEY_CONSTANT)
    return ChainKdfResult(newChainKey, messageKey)
}

/**
 * Derive the initial root key from concatenated X3DH DH outputs.
 * HKDF(ikm=DH1||DH2||DH3[||DH4], salt=32 zeros, info="ShadeX3DH", length=32)
 */
fun deriveInitialRootKey(crypto: CryptoProvider, sharedSecrets: List<ByteArray>): ByteArray {
    val total = sharedSecrets.sumOf { it.size }
    val ikm = ByteArray(total)
    var offset = 0
    for (secret in sharedSecrets) {
        secret.copyInto(ikm, offset)
        offset += secret.size
    }
    return crypto.hkdf(ikm, X3DH_SALT, X3DH_INFO, 32)
}
