package no.zyon.shade.group

import no.zyon.shade.crypto.CryptoProvider
import no.zyon.shade.protocol.kdfChainKey
import java.nio.ByteBuffer

/**
 * Group sender-keys (Sesame). Mirror @shade/core/sender-keys.ts.
 *
 * Each sender maintains a chain key that ratchets forward with `kdfChainKey`
 * — same primitive the Double Ratchet uses for its symmetric chain. Per-message
 * AEAD AAD binds (groupId, senderAddress, iteration) so a captured ciphertext
 * cannot be replayed under a different sender or group:
 *
 *   aad = u16_be(groupIdLen) || groupId || u16_be(senderAddrLen) || senderAddr || u32_be(iteration)
 *
 * Each ciphertext is signed by the sender's Ed25519 key over `aad || ciphertext`,
 * which is what receivers verify before advancing their chain.
 */

data class SenderKeyMessage(
    val senderAddress: String,
    val iteration: Int,
    val ciphertext: ByteArray,
    val nonce: ByteArray,
    val signature: ByteArray,
)

fun encodeSenderHeader(groupId: String, senderAddress: String, iteration: Int): ByteArray {
    val gBytes = groupId.toByteArray(Charsets.UTF_8)
    val sBytes = senderAddress.toByteArray(Charsets.UTF_8)
    require(gBytes.size <= 0xFFFF) { "groupId too long (>65535 UTF-8 bytes)" }
    require(sBytes.size <= 0xFFFF) { "senderAddress too long (>65535 UTF-8 bytes)" }

    val out = ByteArray(2 + gBytes.size + 2 + sBytes.size + 4)
    val buf = ByteBuffer.wrap(out)
    buf.putShort(gBytes.size.toShort())
    buf.put(gBytes)
    buf.putShort(sBytes.size.toShort())
    buf.put(sBytes)
    buf.putInt(iteration)
    return out
}

/**
 * Compute (newChainKey, messageKey, aad) for the next group message.
 * Pure function; caller is responsible for state advancement and the AEAD/sign
 * steps (which need access to the signing private key not exposed here).
 */
data class SenderStepResult(
    val newChainKey: ByteArray,
    val messageKey: ByteArray,
    val aad: ByteArray,
)

fun senderKeyStep(
    crypto: CryptoProvider,
    chainKey: ByteArray,
    groupId: String,
    senderAddress: String,
    iteration: Int,
): SenderStepResult {
    val r = kdfChainKey(crypto, chainKey)
    val aad = encodeSenderHeader(groupId, senderAddress, iteration)
    return SenderStepResult(newChainKey = r.newChainKey, messageKey = r.messageKey, aad = aad)
}

/**
 * Concatenate `aad || ciphertext` — the byte string the sender signs and the
 * receiver verifies. Exposed as a helper so vector parity can pin both sides.
 */
fun senderSignedBytes(aad: ByteArray, ciphertext: ByteArray): ByteArray {
    val out = ByteArray(aad.size + ciphertext.size)
    aad.copyInto(out, 0)
    ciphertext.copyInto(out, aad.size)
    return out
}
