package com.peachpasswords.sync

import org.bouncycastle.crypto.digests.Blake2sDigest
import org.bouncycastle.crypto.generators.HKDFBytesGenerator
import org.bouncycastle.crypto.params.HKDFParameters
import org.bouncycastle.crypto.params.X25519PrivateKeyParameters
import org.bouncycastle.crypto.params.X25519PublicKeyParameters

private const val NOISE_NK_PROTOCOL_NAME = "Noise_NK_25519_ChaChaPoly_BLAKE2s"

interface NoiseNKSession : AutoCloseable {
    fun encrypt(plaintext: ByteArray): ByteArray
    fun decrypt(ciphertext: ByteArray): ByteArray
    val handshakeHash: ByteArray

    fun rekeyOutgoing() {
        throw UnsupportedOperationException("Noise session does not support rekeying")
    }

    fun rekeyIncoming() {
        throw UnsupportedOperationException("Noise session does not support rekeying")
    }

    override fun close() = Unit
}

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

private fun hkdf2(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)
}

// Primitive-only compatibility seams. The production Noise transcript and
// transport state are exclusively owned by the shared Rust core.
internal fun noiseNKInitHashForTest(): ByteArray = initHash()
internal fun noiseNKHkdf2ForTest(
    chainingKey: ByteArray,
    input: ByteArray,
): Pair<ByteArray, ByteArray> = hkdf2(chainingKey, input)

internal fun x25519SharedKey(privateKey: ByteArray, publicKey: ByteArray): ByteArray {
    require(privateKey.size == 32 && publicKey.size == 32) { "X25519 keys must be 32 bytes" }
    val private = X25519PrivateKeyParameters(privateKey, 0)
    val public = X25519PublicKeyParameters(publicKey, 0)
    val output = ByteArray(32)
    private.generateSecret(public, output, 0)
    require(output.any { it.toInt() != 0 }) { "X25519 low-order point rejected" }
    return output
}

suspend fun noiseNKInitiator(
    responderStaticPublicKey: ByteArray,
    send: suspend (ByteArray) -> Unit,
    receive: suspend () -> ByteArray,
    ephemeralPrivateKeyForTest: ByteArray? = null,
): NoiseNKSession {
    val native = RustNoiseSession(
        if (ephemeralPrivateKeyForTest == null) {
            RustNoiseNative.nativeCreateNkInitiator(responderStaticPublicKey)
        } else {
            RustNoiseNative.nativeCreateNkInitiatorForTest(
                responderStaticPublicKey,
                ephemeralPrivateKeyForTest,
            )
        },
    )
    return try {
        send(native.writeHandshake())
        native.readHandshake(receive())
        native
    } catch (error: Throwable) {
        native.close()
        throw error
    }
}

suspend fun noiseNKResponder(
    ourStaticPrivateKey: ByteArray,
    ourStaticPublicKey: ByteArray,
    send: suspend (ByteArray) -> Unit,
    receive: suspend () -> ByteArray,
    ephemeralPrivateKeyForTest: ByteArray? = null,
): NoiseNKSession {
    val native = RustNoiseSession(
        if (ephemeralPrivateKeyForTest == null) {
            RustNoiseNative.nativeCreateNkResponder(
                ourStaticPrivateKey,
                ourStaticPublicKey,
            )
        } else {
            RustNoiseNative.nativeCreateNkResponderForTest(
                ourStaticPrivateKey,
                ourStaticPublicKey,
                ephemeralPrivateKeyForTest,
            )
        },
    )
    return try {
        native.readHandshake(receive())
        send(native.writeHandshake())
        native
    } catch (error: Throwable) {
        native.close()
        throw error
    }
}

fun derivePairingSAS(handshakeHash: ByteArray): String =
    SASDerivation.derive(handshakeHash)