package com.peachpasswords.sync

import org.bouncycastle.crypto.digests.Blake2sDigest
import org.bouncycastle.crypto.generators.HKDFBytesGenerator
import org.bouncycastle.crypto.modes.ChaCha20Poly1305
import org.bouncycastle.crypto.params.AEADParameters
import org.bouncycastle.crypto.params.HKDFParameters
import org.bouncycastle.crypto.params.KeyParameter
import org.bouncycastle.crypto.params.X25519PrivateKeyParameters
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.security.SecureRandom

private const val NOISE_KK_PROTOCOL_NAME = "Noise_KK_25519_ChaChaPoly_BLAKE2s"

internal data class NoiseKKKeyPair(
    val publicKey: ByteArray,
    val secretKey: ByteArray,
)

private fun initKKHash(): ByteArray {
    val name = NOISE_KK_PROTOCOL_NAME.encodeToByteArray()
    if (name.size <= 32) return ByteArray(32).also { name.copyInto(it) }
    return Blake2sDigest(256).run {
        update(name, 0, name.size)
        ByteArray(32).also { doFinal(it, 0) }
    }
}

private fun mixKKHash(hash: ByteArray, data: ByteArray): ByteArray =
    Blake2sDigest(256).run {
        update(hash, 0, hash.size)
        update(data, 0, data.size)
        ByteArray(32).also { doFinal(it, 0) }
    }

private fun kkHkdf2(chainingKey: ByteArray, input: ByteArray): Pair<ByteArray, ByteArray> {
    val generator = HKDFBytesGenerator(Blake2sDigest(256))
    generator.init(HKDFParameters(input, chainingKey, ByteArray(0)))
    val output = ByteArray(64)
    generator.generateBytes(output, 0, output.size)
    return output.copyOfRange(0, 32) to output.copyOfRange(32, 64)
}

private fun kkNonce(value: Long): ByteArray =
    ByteArray(12).also {
        ByteBuffer.wrap(it).order(ByteOrder.LITTLE_ENDIAN).putLong(4, value)
    }

private fun kkSeal(key: ByteArray, plaintext: ByteArray, aad: ByteArray): ByteArray {
    val cipher = ChaCha20Poly1305()
    cipher.init(true, AEADParameters(KeyParameter(key), 128, ByteArray(12), aad))
    val output = ByteArray(cipher.getOutputSize(plaintext.size))
    var size = cipher.processBytes(plaintext, 0, plaintext.size, output, 0)
    size += cipher.doFinal(output, size)
    return output.copyOf(size)
}

private fun kkOpen(key: ByteArray, ciphertext: ByteArray, aad: ByteArray): ByteArray? =
    runCatching {
        val cipher = ChaCha20Poly1305()
        cipher.init(false, AEADParameters(KeyParameter(key), 128, ByteArray(12), aad))
        val output = ByteArray(cipher.getOutputSize(ciphertext.size))
        var size = cipher.processBytes(ciphertext, 0, ciphertext.size, output, 0)
        size += cipher.doFinal(output, size)
        output.copyOf(size)
    }.getOrNull()

private fun freshKKKeyPair(): NoiseKKKeyPair {
    val privateKey = X25519PrivateKeyParameters(SecureRandom())
    return NoiseKKKeyPair(privateKey.generatePublicKey().encoded, privateKey.encoded)
}

/**
 * Standard Noise KK reconnect initiator.
 * Pre-messages: -> s, <- s; messages: -> e,es,ss then <- e,ee,se.
 */
internal suspend fun noiseKKInitiator(
    ourStaticPrivateKey: ByteArray,
    ourStaticPublicKey: ByteArray,
    responderStaticPublicKey: ByteArray,
    send: suspend (ByteArray) -> Unit,
    receive: suspend () -> ByteArray,
    ephemeral: NoiseKKKeyPair = freshKKKeyPair(),
): NoiseNKSession {
    requireKKKeys(
        ourStaticPrivateKey,
        ourStaticPublicKey,
        responderStaticPublicKey,
        ephemeral.secretKey,
        ephemeral.publicKey,
    )
    var hash = initKKHash()
    var chainingKey = hash.copyOf()
    hash = mixKKHash(hash, ourStaticPublicKey)
    hash = mixKKHash(hash, responderStaticPublicKey)

    // -> e, es, ss
    hash = mixKKHash(hash, ephemeral.publicKey)
    var keys = kkHkdf2(
        chainingKey,
        x25519SharedKey(ephemeral.secretKey, responderStaticPublicKey),
    )
    chainingKey = keys.first
    keys = kkHkdf2(
        chainingKey,
        x25519SharedKey(ourStaticPrivateKey, responderStaticPublicKey),
    )
    chainingKey = keys.first
    val encrypted1 = kkSeal(keys.second, ByteArray(0), hash)
    hash = mixKKHash(hash, encrypted1)
    send(ephemeral.publicKey + encrypted1)

    // <- e, ee, se
    val message2 = receive()
    require(message2.size >= 48) { "Noise KK: msg2 too short" }
    val responderEphemeral = message2.copyOfRange(0, 32)
    val encrypted2 = message2.copyOfRange(32, message2.size)
    hash = mixKKHash(hash, responderEphemeral)
    keys = kkHkdf2(
        chainingKey,
        x25519SharedKey(ephemeral.secretKey, responderEphemeral),
    )
    chainingKey = keys.first
    keys = kkHkdf2(
        chainingKey,
        x25519SharedKey(ourStaticPrivateKey, responderEphemeral),
    )
    chainingKey = keys.first
    requireNotNull(kkOpen(keys.second, encrypted2, hash)) {
        "Noise KK: msg2 auth failed — wrong peer key or MITM"
    }
    hash = mixKKHash(hash, encrypted2)

    val transportKeys = kkHkdf2(chainingKey, ByteArray(0))
    return kkTransportSession(transportKeys.first, transportKeys.second, hash)
}

/**
 * Standard Noise KK responder. [expectedInitiatorStaticPublicKey] must be
 * selected from local paired state before the handshake; it is never learned
 * from a peer capability or accepted via fallback.
 */
internal suspend fun noiseKKResponder(
    ourStaticPrivateKey: ByteArray,
    ourStaticPublicKey: ByteArray,
    expectedInitiatorStaticPublicKey: ByteArray,
    send: suspend (ByteArray) -> Unit,
    receive: suspend () -> ByteArray,
    ephemeral: NoiseKKKeyPair = freshKKKeyPair(),
): NoiseNKSession {
    requireKKKeys(
        ourStaticPrivateKey,
        ourStaticPublicKey,
        expectedInitiatorStaticPublicKey,
        ephemeral.secretKey,
        ephemeral.publicKey,
    )
    var hash = initKKHash()
    var chainingKey = hash.copyOf()
    hash = mixKKHash(hash, expectedInitiatorStaticPublicKey)
    hash = mixKKHash(hash, ourStaticPublicKey)

    // -> e, es, ss
    val message1 = receive()
    require(message1.size >= 48) { "Noise KK: msg1 too short" }
    val initiatorEphemeral = message1.copyOfRange(0, 32)
    val encrypted1 = message1.copyOfRange(32, message1.size)
    hash = mixKKHash(hash, initiatorEphemeral)
    var keys = kkHkdf2(
        chainingKey,
        x25519SharedKey(ourStaticPrivateKey, initiatorEphemeral),
    )
    chainingKey = keys.first
    keys = kkHkdf2(
        chainingKey,
        x25519SharedKey(ourStaticPrivateKey, expectedInitiatorStaticPublicKey),
    )
    chainingKey = keys.first
    requireNotNull(kkOpen(keys.second, encrypted1, hash)) {
        "Noise KK: msg1 auth failed — wrong peer key or MITM"
    }
    hash = mixKKHash(hash, encrypted1)

    // <- e, ee, se
    hash = mixKKHash(hash, ephemeral.publicKey)
    keys = kkHkdf2(
        chainingKey,
        x25519SharedKey(ephemeral.secretKey, initiatorEphemeral),
    )
    chainingKey = keys.first
    keys = kkHkdf2(
        chainingKey,
        x25519SharedKey(ephemeral.secretKey, expectedInitiatorStaticPublicKey),
    )
    chainingKey = keys.first
    val encrypted2 = kkSeal(keys.second, ByteArray(0), hash)
    hash = mixKKHash(hash, encrypted2)
    send(ephemeral.publicKey + encrypted2)

    val transportKeys = kkHkdf2(chainingKey, ByteArray(0))
    return kkTransportSession(transportKeys.second, transportKeys.first, hash)
}

private fun requireKKKeys(vararg keys: ByteArray) {
    require(keys.all { it.size == 32 }) { "Noise KK keys must be 32 bytes" }
}

private fun kkTransportSession(
    sendKey: ByteArray,
    receiveKey: ByteArray,
    hash: ByteArray,
): NoiseNKSession {
    var sendNonce = 0L
    var receiveNonce = 0L
    return object : NoiseNKSession {
        override fun encrypt(plaintext: ByteArray): ByteArray {
            val cipher = ChaCha20Poly1305()
            cipher.init(true, AEADParameters(KeyParameter(sendKey), 128, kkNonce(sendNonce++), ByteArray(0)))
            val output = ByteArray(cipher.getOutputSize(plaintext.size))
            var size = cipher.processBytes(plaintext, 0, plaintext.size, output, 0)
            size += cipher.doFinal(output, size)
            return output.copyOf(size)
        }

        override fun decrypt(ciphertext: ByteArray): ByteArray {
            val cipher = ChaCha20Poly1305()
            cipher.init(false, AEADParameters(KeyParameter(receiveKey), 128, kkNonce(receiveNonce++), ByteArray(0)))
            val output = ByteArray(cipher.getOutputSize(ciphertext.size))
            var size = cipher.processBytes(ciphertext, 0, ciphertext.size, output, 0)
            size += cipher.doFinal(output, size)
            return output.copyOf(size)
        }

        override val handshakeHash: ByteArray = hash.copyOf()
    }
}