package no.zyon.shade.protocol

import no.zyon.shade.crypto.CryptoProvider
import no.zyon.shade.types.ChainState
import no.zyon.shade.types.Constants
import no.zyon.shade.types.KeyPair
import no.zyon.shade.types.RatchetMessage
import no.zyon.shade.types.SessionState
import java.nio.ByteBuffer

/**
 * Double Ratchet implementation. Mirrors @shade/core/ratchet.ts.
 *
 * Must produce byte-identical ciphertext to the TypeScript version
 * for the same inputs.
 */

// ─── Session initialization ─────────────────────────────────

fun initSenderSession(
    crypto: CryptoProvider,
    rootKey: ByteArray,
    remoteIdentityKey: ByteArray,
    remoteDHPublicKey: ByteArray,
): SessionState {
    val (dhSendPub, dhSendPriv) = crypto.generateX25519KeyPair()
    val dhOutput = crypto.x25519(dhSendPriv, remoteDHPublicKey)
    val (newRootKey, chainKey) = kdfRootKey(crypto, rootKey, dhOutput).let {
        it.newRootKey to it.chainKey
    }
    return SessionState(
        remoteIdentityKey = remoteIdentityKey,
        rootKey = newRootKey,
        sendChain = ChainState(chainKey = chainKey, counter = 0),
        receiveChain = null,
        dhSend = KeyPair(publicKey = dhSendPub, privateKey = dhSendPriv),
        dhReceive = remoteDHPublicKey,
        previousSendCounter = 0,
        skippedKeys = mutableMapOf(),
    )
}

fun initReceiverSession(
    rootKey: ByteArray,
    remoteIdentityKey: ByteArray,
    localDHKeyPair: KeyPair,
): SessionState {
    return SessionState(
        remoteIdentityKey = remoteIdentityKey,
        rootKey = rootKey,
        sendChain = ChainState(chainKey = ByteArray(32), counter = 0),
        receiveChain = null,
        dhSend = localDHKeyPair,
        dhReceive = null,
        previousSendCounter = 0,
        skippedKeys = mutableMapOf(),
    )
}

// ─── Header encoding (for AES-GCM AAD) ──────────────────────

private fun encodeHeader(
    dhPublicKey: ByteArray,
    previousCounter: Int,
    counter: Int,
): ByteArray {
    val buf = ByteBuffer.allocate(40)
    buf.put(dhPublicKey)
    buf.putInt(previousCounter) // big-endian by default in ByteBuffer
    buf.putInt(counter)
    return buf.array()
}

// ─── Encrypt ─────────────────────────────────────────────────

fun ratchetEncrypt(
    crypto: CryptoProvider,
    session: SessionState,
    plaintext: ByteArray,
): RatchetMessage {
    val oldChainKey = session.sendChain.chainKey
    val (newChainKey, messageKey) = kdfChainKey(crypto, oldChainKey).let {
        it.newChainKey to it.messageKey
    }
    crypto.zeroize(oldChainKey)

    val counter = session.sendChain.counter
    val header = encodeHeader(session.dhSend.publicKey, session.previousSendCounter, counter)

    val (ciphertext, nonce) = crypto.aesGcmEncrypt(messageKey, plaintext, header)
    crypto.zeroize(messageKey)

    session.sendChain.chainKey = newChainKey
    session.sendChain.counter = counter + 1

    return RatchetMessage(
        dhPublicKey = session.dhSend.publicKey,
        previousCounter = session.previousSendCounter,
        counter = counter,
        ciphertext = ciphertext,
        nonce = nonce,
    )
}

// ─── Decrypt ─────────────────────────────────────────────────

private fun skippedKeyId(dhPublicKey: ByteArray, counter: Int): String {
    return dhPublicKey.joinToString("") { "%02x".format(it) } + ":" + counter
}

fun ratchetDecrypt(
    crypto: CryptoProvider,
    session: SessionState,
    message: RatchetMessage,
): ByteArray {
    // Case 1: skipped key
    val skipId = skippedKeyId(message.dhPublicKey, message.counter)
    val skippedKey = session.skippedKeys[skipId]
    if (skippedKey != null) {
        session.skippedKeys.remove(skipId)
        try {
            return decryptWithKey(crypto, skippedKey, message)
        } finally {
            crypto.zeroize(skippedKey)
        }
    }

    // Case 2 or 3: DH ratchet check
    val isNewRatchet = session.dhReceive == null ||
        !message.dhPublicKey.contentEquals(session.dhReceive!!)

    if (isNewRatchet) {
        if (session.receiveChain != null && session.dhReceive != null) {
            skipMessageKeys(
                crypto,
                session,
                session.dhReceive!!,
                session.receiveChain!!,
                message.previousCounter,
            )
        }
        performDHRatchetStep(crypto, session, message.dhPublicKey)
    }

    val receiveChain = session.receiveChain
        ?: throw IllegalStateException("No receiving chain available")

    skipMessageKeys(crypto, session, message.dhPublicKey, receiveChain, message.counter)

    val oldChainKey = receiveChain.chainKey
    val (newChainKey, messageKey) = kdfChainKey(crypto, oldChainKey).let {
        it.newChainKey to it.messageKey
    }
    crypto.zeroize(oldChainKey)
    receiveChain.chainKey = newChainKey
    receiveChain.counter = message.counter + 1

    try {
        return decryptWithKey(crypto, messageKey, message)
    } finally {
        crypto.zeroize(messageKey)
    }
}

private fun performDHRatchetStep(
    crypto: CryptoProvider,
    session: SessionState,
    remoteDHKey: ByteArray,
) {
    session.previousSendCounter = session.sendChain.counter
    session.dhReceive = remoteDHKey

    // DH with current send key → new receiving chain
    val dh1 = crypto.x25519(session.dhSend.privateKey, remoteDHKey)
    val oldRootKey1 = session.rootKey
    val recv = kdfRootKey(crypto, oldRootKey1, dh1)
    crypto.zeroize(oldRootKey1)
    crypto.zeroize(dh1)
    session.rootKey = recv.newRootKey
    session.receiveChain = ChainState(chainKey = recv.chainKey, counter = 0)

    // Generate new DH keypair, zero old private
    val oldDhPrivate = session.dhSend.privateKey
    val (newDhPub, newDhPriv) = crypto.generateX25519KeyPair()
    session.dhSend = KeyPair(publicKey = newDhPub, privateKey = newDhPriv)
    crypto.zeroize(oldDhPrivate)

    // DH with new send key → new sending chain
    val dh2 = crypto.x25519(newDhPriv, remoteDHKey)
    val oldRootKey2 = session.rootKey
    val send = kdfRootKey(crypto, oldRootKey2, dh2)
    crypto.zeroize(oldRootKey2)
    crypto.zeroize(dh2)
    session.rootKey = send.newRootKey
    if (session.sendChain.chainKey.isNotEmpty()) {
        crypto.zeroize(session.sendChain.chainKey)
    }
    session.sendChain = ChainState(chainKey = send.chainKey, counter = 0)
}

private fun skipMessageKeys(
    crypto: CryptoProvider,
    session: SessionState,
    dhPublicKey: ByteArray,
    chain: ChainState,
    untilCounter: Int,
) {
    val toSkip = untilCounter - chain.counter
    if (toSkip < 0) return
    if (toSkip > Constants.MAX_SKIP) {
        throw IllegalStateException("Cannot skip $toSkip messages (max: ${Constants.MAX_SKIP})")
    }

    for (i in chain.counter until untilCounter) {
        val (newChainKey, messageKey) = kdfChainKey(crypto, chain.chainKey).let {
            it.newChainKey to it.messageKey
        }
        val id = skippedKeyId(dhPublicKey, i)
        session.skippedKeys[id] = messageKey
        chain.chainKey = newChainKey
        chain.counter = i + 1

        while (session.skippedKeys.size > Constants.MAX_CACHED_SKIPPED_KEYS) {
            val firstKey = session.skippedKeys.keys.first()
            session.skippedKeys.remove(firstKey)
        }
    }
}

private fun decryptWithKey(
    crypto: CryptoProvider,
    messageKey: ByteArray,
    message: RatchetMessage,
): ByteArray {
    val aad = encodeHeader(message.dhPublicKey, message.previousCounter, message.counter)
    return crypto.aesGcmDecrypt(messageKey, message.ciphertext, message.nonce, aad)
}
