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 org.bouncycastle.crypto.params.X25519PublicKeyParameters
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.security.SecureRandom

private const val NOISE_NK_PROTOCOL_NAME = "Noise_NK_25519_ChaChaPoly_BLAKE2s"

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

private data class NoiseNKKeyPair(
    val publicKey: ByteArray,
    val secretKey: ByteArray,
)

private fun initHash(): ByteArray {
    val nameBytes = NOISE_NK_PROTOCOL_NAME.encodeToByteArray()
    if (nameBytes.size <= 32) {
        return ByteArray(32).also { nameBytes.copyInto(it) }
    }
    // PCH-01-002: the named pattern hash is BLAKE2s; the protocol-name init hash
    // MUST use it too (Noise spec §5.2), not SHA-256. Coordinated wire change with
    // peach-browser noise-js.ts.
    val blake = Blake2sDigest(256)
    blake.update(nameBytes, 0, nameBytes.size)
    val out = ByteArray(32)
    blake.doFinal(out, 0)
    return out
}

private fun mixHash(h: ByteArray, data: ByteArray): ByteArray {
    val blake = Blake2sDigest(256)
    blake.update(h, 0, h.size)
    blake.update(data, 0, data.size)
    val out = ByteArray(32)
    blake.doFinal(out, 0)
    return out
}

private fun hkdf2(ck: ByteArray, ikm: ByteArray): Pair<ByteArray, ByteArray> {
    // PCH-01-002: HKDF must run over the named hash (BLAKE2s), not SHA-256.
    val generator = HKDFBytesGenerator(Blake2sDigest(256))
    generator.init(HKDFParameters(ikm, ck, ByteArray(0)))
    val out = ByteArray(64)
    generator.generateBytes(out, 0, out.size)
    return out.copyOfRange(0, 32) to out.copyOfRange(32, 64)
}

// PCH-01-002: test-only seams for the cross-platform BLAKE2s vectors pinned in
// CrossPlatformKATTest.kt (must stay byte-identical to peach-browser noise-js.ts).
internal fun noiseNKInitHashForTest(): ByteArray = initHash()
internal fun noiseNKHkdf2ForTest(ck: ByteArray, ikm: ByteArray): Pair<ByteArray, ByteArray> = hkdf2(ck, ikm)

private fun nonceBytes(n: Long): ByteArray {
    val buf = ByteArray(12)
    ByteBuffer.wrap(buf).order(ByteOrder.LITTLE_ENDIAN).putLong(4, n)
    return buf
}

private fun generateKeyPair(): NoiseNKKeyPair {
    val priv = X25519PrivateKeyParameters(SecureRandom())
    return NoiseNKKeyPair(
        publicKey = priv.generatePublicKey().encoded,
        secretKey = priv.encoded,
    )
}

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

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

private fun chachaOpen(key: ByteArray, nonce: ByteArray, ciphertext: ByteArray, aad: ByteArray): ByteArray? {
    return try {
        val cipher = ChaCha20Poly1305()
        cipher.init(false, AEADParameters(KeyParameter(key), 128, nonce, aad))
        val out = ByteArray(cipher.getOutputSize(ciphertext.size))
        var len = cipher.processBytes(ciphertext, 0, ciphertext.size, out, 0)
        len += cipher.doFinal(out, len)
        out.copyOf(len)
    } catch (_: Throwable) {
        null
    }
}

suspend fun noiseNKInitiator(
    responderStaticPublicKey: ByteArray,
    send: suspend (ByteArray) -> Unit,
    receive: suspend () -> ByteArray,
): NoiseNKSession {
    var h = initHash()
    var ck = h.copyOf()

    h = mixHash(h, responderStaticPublicKey)

    val ephemeral = generateKeyPair()
    h = mixHash(h, ephemeral.publicKey)
    val es = x25519SharedKey(ephemeral.secretKey, responderStaticPublicKey)
    val (ck1, k1) = hkdf2(ck, es)
    ck = ck1

    val cipher1 = ChaCha20Poly1305()
    cipher1.init(true, AEADParameters(KeyParameter(k1), 128, nonceBytes(0L), h))
    val encrypted1 = ByteArray(cipher1.getOutputSize(0))
    cipher1.doFinal(encrypted1, 0)
    h = mixHash(h, encrypted1)

    val msg1 = ByteArray(32 + encrypted1.size)
    ephemeral.publicKey.copyInto(msg1, 0)
    encrypted1.copyInto(msg1, 32)
    send(msg1)

    val msg2 = receive()
    require(msg2.size >= 48) { "Noise NK: msg2 too short" }
    val remoteEphemeralPublic = msg2.copyOfRange(0, 32)
    val encrypted2 = msg2.copyOfRange(32, msg2.size)

    h = mixHash(h, remoteEphemeralPublic)
    val ee = x25519SharedKey(ephemeral.secretKey, remoteEphemeralPublic)
    val (ck2, k2) = hkdf2(ck, ee)
    ck = ck2

    val decrypted2 = chachaOpen(k2, nonceBytes(0L), encrypted2, h)
    requireNotNull(decrypted2) { "Noise NK: msg2 auth failed — wrong responder key or MITM" }
    h = mixHash(h, encrypted2)

    val (sendKey, recvKey) = hkdf2(ck, ByteArray(0))
    val sessionHash = h.copyOf()

    var sendNonce = 0L
    var recvNonce = 0L

    return object : NoiseNKSession {
        override fun encrypt(plaintext: ByteArray): ByteArray {
            return chachaSeal(sendKey, nonceBytes(sendNonce++), plaintext, ByteArray(0))
        }

        override fun decrypt(ciphertext: ByteArray): ByteArray {
            return chachaOpen(recvKey, nonceBytes(recvNonce++), ciphertext, ByteArray(0))
                ?: throw IllegalStateException("Noise NK: decryption failed")
        }

        override val handshakeHash: ByteArray = sessionHash
    }
}

suspend fun noiseNKResponder(
    ourStaticPrivateKey: ByteArray,
    ourStaticPublicKey: ByteArray,
    send: suspend (ByteArray) -> Unit,
    receive: suspend () -> ByteArray,
): NoiseNKSession {
    var h = initHash()
    var ck = h.copyOf()

    h = mixHash(h, ourStaticPublicKey)

    val msg1 = receive()
    require(msg1.size >= 48) { "Noise NK: msg1 too short" }
    val remoteEphemeralPublic = msg1.copyOfRange(0, 32)
    val encrypted1 = msg1.copyOfRange(32, msg1.size)

    h = mixHash(h, remoteEphemeralPublic)
    val es = x25519SharedKey(ourStaticPrivateKey, remoteEphemeralPublic)
    val (ck1, k1) = hkdf2(ck, es)
    ck = ck1

    val decrypted1 = chachaOpen(k1, nonceBytes(0L), encrypted1, h)
    requireNotNull(decrypted1) { "Noise NK: msg1 auth failed — wrong static key or MITM" }
    h = mixHash(h, encrypted1)

    val ephemeral = generateKeyPair()
    h = mixHash(h, ephemeral.publicKey)
    val ee = x25519SharedKey(ephemeral.secretKey, remoteEphemeralPublic)
    val (ck2, k2) = hkdf2(ck, ee)
    ck = ck2

    val encrypted2 = chachaSeal(k2, nonceBytes(0L), ByteArray(0), h)
    h = mixHash(h, encrypted2)

    val msg2 = ByteArray(32 + encrypted2.size)
    ephemeral.publicKey.copyInto(msg2, 0)
    encrypted2.copyInto(msg2, 32)
    send(msg2)

    val (initiatorSendKey, responderSendKey) = hkdf2(ck, ByteArray(0))
    val sessionHash = h.copyOf()

    var sendNonce = 0L
    var recvNonce = 0L

    return object : NoiseNKSession {
        override fun encrypt(plaintext: ByteArray): ByteArray {
            return chachaSeal(responderSendKey, nonceBytes(sendNonce++), plaintext, ByteArray(0))
        }

        override fun decrypt(ciphertext: ByteArray): ByteArray {
            return chachaOpen(initiatorSendKey, nonceBytes(recvNonce++), ciphertext, ByteArray(0))
                ?: throw IllegalStateException("Noise NK: decryption failed")
        }

        override val handshakeHash: ByteArray = sessionHash
    }
}

fun derivePairingSAS(handshakeHash: ByteArray): String {
    return SASDerivation.derive(handshakeHash)
}