package com.peachpasswords.sync

import org.bouncycastle.crypto.InvalidCipherTextException
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 PROTOCOL_NAME = "Noise_XX_25519_ChaChaPoly_BLAKE2s"

private fun initHash(): ByteArray {
    val nameBytes = 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-xx.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-xx.ts).
internal fun noiseXXInitHashForTest(): ByteArray = initHash()
internal fun noiseXXHkdf2ForTest(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
}

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

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

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

private fun keyPairFromSeed(seed: ByteArray): X25519KeyPair {
    val priv = X25519PrivateKeyParameters(seed, 0)
    return X25519KeyPair(
        publicKey = priv.generatePublicKey().encoded,
        secretKey = priv.encoded,
    )
}

private fun sharedKey(privateKey: ByteArray, publicKey: ByteArray): ByteArray {
    val priv = X25519PrivateKeyParameters(privateKey, 0)
    val pub = X25519PublicKeyParameters(publicKey, 0)
    val out = ByteArray(32)
    priv.generateSecret(pub, out, 0)
    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 (_: InvalidCipherTextException) {
        null
    }
}

internal fun noiseXXDeterministicHandshakeHashForTest(): ByteArray {
    val initiatorStatic = keyPairFromSeed(ByteArray(32) { 0xa1.toByte() })
    val responderStatic = keyPairFromSeed(ByteArray(32) { 0xb2.toByte() })
    val initiatorEphemeral = keyPairFromSeed(ByteArray(32) { 0x01 })
    val responderEphemeral = keyPairFromSeed(ByteArray(32) { 0x02 })

    var hI = initHash()
    var ckI = hI.copyOf()
    var hR = initHash()
    var ckR = hR.copyOf()

    // Message 1: -> e
    val msg1 = initiatorEphemeral.publicKey
    hI = mixHash(hI, msg1)
    hI = mixHash(hI, ByteArray(0))
    hR = mixHash(hR, msg1)
    hR = mixHash(hR, ByteArray(0))

    // Message 2: <- e, ee, s, es
    hR = mixHash(hR, responderEphemeral.publicKey)
    val eeR = sharedKey(responderEphemeral.secretKey, msg1)
    val (ckR1, kR1) = hkdf2(ckR, eeR)
    ckR = ckR1
    val encryptedResponderStatic = chachaSeal(kR1, nonceBytes(0L), responderStatic.publicKey, hR)
    hR = mixHash(hR, encryptedResponderStatic)

    val esR = sharedKey(responderStatic.secretKey, msg1)
    val (ckR2, kR2) = hkdf2(ckR, esR)
    ckR = ckR2
    val encryptedPayload = chachaSeal(kR2, nonceBytes(0L), ByteArray(0), hR)
    hR = mixHash(hR, encryptedPayload)

    val remoteEphemeral = responderEphemeral.publicKey
    hI = mixHash(hI, remoteEphemeral)
    val eeI = sharedKey(initiatorEphemeral.secretKey, remoteEphemeral)
    val (ckI1, kI1) = hkdf2(ckI, eeI)
    ckI = ckI1
    val remoteStaticPlain = chachaOpen(kI1, nonceBytes(0L), encryptedResponderStatic, hI)
        ?: throw IllegalStateException("Noise XX KAT: msg2 static auth failed")
    hI = mixHash(hI, encryptedResponderStatic)

    val esI = sharedKey(initiatorEphemeral.secretKey, remoteStaticPlain)
    val (ckI2, kI2) = hkdf2(ckI, esI)
    ckI = ckI2
    chachaOpen(kI2, nonceBytes(0L), encryptedPayload, hI)
        ?: throw IllegalStateException("Noise XX KAT: msg2 payload auth failed")
    hI = mixHash(hI, encryptedPayload)

    // Message 3: -> s, se
    val encryptedInitiatorStatic = chachaSeal(kI2, nonceBytes(1L), initiatorStatic.publicKey, hI)
    hI = mixHash(hI, encryptedInitiatorStatic)
    val seI = sharedKey(initiatorStatic.secretKey, remoteEphemeral)
    val (ckI3, kI3) = hkdf2(ckI, seI)
    ckI = ckI3
    val encryptedFinalPayload = chachaSeal(kI3, nonceBytes(0L), ByteArray(0), hI)
    hI = mixHash(hI, encryptedFinalPayload)

    val initiatorStaticPlain = chachaOpen(kR2, nonceBytes(1L), encryptedInitiatorStatic, hR)
        ?: throw IllegalStateException("Noise XX KAT: msg3 static auth failed")
    hR = mixHash(hR, encryptedInitiatorStatic)
    val seR = sharedKey(responderEphemeral.secretKey, initiatorStaticPlain)
    val (ckR3, kR3) = hkdf2(ckR, seR)
    ckR = ckR3
    chachaOpen(kR3, nonceBytes(0L), encryptedFinalPayload, hR)
        ?: throw IllegalStateException("Noise XX KAT: msg3 final payload auth failed")
    hR = mixHash(hR, encryptedFinalPayload)

    check(hI.contentEquals(hR)) { "Noise XX KAT sides diverged" }
    check(ckI.contentEquals(ckR)) { "Noise XX KAT chaining keys diverged" }
    return hI
}

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

    val e = generateKeyPair()
    h = mixHash(h, e.publicKey)
    // PCH-02-002: canonical Noise still runs EncryptAndHash(empty) for msg1.
    // No bytes are added to the wire; the transcript hash moves in lockstep
    // with peach-browser.
    h = mixHash(h, ByteArray(0))
    send(e.publicKey)

    val msg2 = receive()
    require(msg2.size >= 32 + 32 + 16 + 16) { "Noise XX: msg2 too short" }
    val remoteEphemeral = msg2.copyOfRange(0, 32)
    val encryptedStatic = msg2.copyOfRange(32, 32 + 32 + 16)
    val encryptedPayload = msg2.copyOfRange(32 + 32 + 16, msg2.size)

    h = mixHash(h, remoteEphemeral)
    val ee = sharedKey(e.secretKey, remoteEphemeral)
    val (ck1, k1) = hkdf2(ck, ee)
    ck = ck1

    val remoteStaticPlain = chachaOpen(k1, nonceBytes(0L), encryptedStatic, h)
        ?: throw IllegalStateException("Noise XX: msg2 static auth failed")
    require(remoteStaticPlain.size == 32) { "Noise XX: msg2 static auth failed" }
    h = mixHash(h, encryptedStatic)

    val es = sharedKey(e.secretKey, remoteStaticPlain)
    val (ck2, k2) = hkdf2(ck, es)
    ck = ck2

    val payloadPlain = chachaOpen(k2, nonceBytes(0L), encryptedPayload, h)
        ?: throw IllegalStateException("Noise XX: msg2 payload auth failed")
    h = mixHash(h, encryptedPayload)
    @Suppress("UNUSED_VARIABLE")
    val _payloadPlain = payloadPlain

    val encryptedOurStatic = chachaSeal(k2, nonceBytes(1L), ourStaticPublicKey, h)
    h = mixHash(h, encryptedOurStatic)

    val se = sharedKey(ourStaticPrivateKey, remoteEphemeral)
    val (ck3, k3) = hkdf2(ck, se)
    ck = ck3

    val encryptedFinalPayload = chachaSeal(k3, nonceBytes(0L), ByteArray(0), h)
    h = mixHash(h, encryptedFinalPayload)

    val msg3 = ByteArray(encryptedOurStatic.size + encryptedFinalPayload.size)
    encryptedOurStatic.copyInto(msg3, 0)
    encryptedFinalPayload.copyInto(msg3, encryptedOurStatic.size)
    send(msg3)

    val (initSend, respSend) = hkdf2(ck, ByteArray(0))
    val sendCipher = initSend
    val recvCipher = respSend
    var sendNonce = 0L
    var recvNonce = 0L

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

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

        override val remoteStaticPublicKey: ByteArray = remoteStaticPlain
        override val handshakeHash: ByteArray = h.copyOf()
    }
}

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

    val msg1 = receive()
    require(msg1.size >= 32) { "Noise XX: msg1 too short" }
    val remoteEphemeral = msg1.copyOfRange(0, 32)
    h = mixHash(h, remoteEphemeral)
    // PCH-02-002: mirror the initiator's empty-payload MixHash for msg1.
    h = mixHash(h, ByteArray(0))

    val e = generateKeyPair()
    h = mixHash(h, e.publicKey)
    val ee = sharedKey(e.secretKey, remoteEphemeral)
    val (ck1, k1) = hkdf2(ck, ee)
    ck = ck1

    val encryptedOurStatic = chachaSeal(k1, nonceBytes(0L), ourStaticPublicKey, h)
    h = mixHash(h, encryptedOurStatic)

    val es = sharedKey(ourStaticPrivateKey, remoteEphemeral)
    val (ck2, k2) = hkdf2(ck, es)
    ck = ck2

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

    val msg2 = ByteArray(e.publicKey.size + encryptedOurStatic.size + encryptedPayload.size)
    e.publicKey.copyInto(msg2, 0)
    encryptedOurStatic.copyInto(msg2, e.publicKey.size)
    encryptedPayload.copyInto(msg2, e.publicKey.size + encryptedOurStatic.size)
    send(msg2)

    val msg3 = receive()
    require(msg3.size >= 32 + 16 + 16) { "Noise XX: msg3 too short" }
    val encryptedRemoteStatic = msg3.copyOfRange(0, 32 + 16)
    val encryptedFinalPayload = msg3.copyOfRange(32 + 16, msg3.size)

    val remoteStaticPlain = chachaOpen(k2, nonceBytes(1L), encryptedRemoteStatic, h)
        ?: throw IllegalStateException("Noise XX: msg3 static auth failed")
    require(remoteStaticPlain.size == 32) { "Noise XX: msg3 static auth failed" }
    h = mixHash(h, encryptedRemoteStatic)

    val se = sharedKey(e.secretKey, remoteStaticPlain)
    val (ck3, k3) = hkdf2(ck, se)
    ck = ck3

    val finalPayloadPlain = chachaOpen(k3, nonceBytes(0L), encryptedFinalPayload, h)
        ?: throw IllegalStateException("Noise XX: msg3 final payload auth failed")
    h = mixHash(h, encryptedFinalPayload)
    @Suppress("UNUSED_VARIABLE")
    val _finalPayloadPlain = finalPayloadPlain

    val (initSend, respSend) = hkdf2(ck, ByteArray(0))
    val sendCipher = respSend
    val recvCipher = initSend
    var sendNonce = 0L
    var recvNonce = 0L

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

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

        override val remoteStaticPublicKey: ByteArray = remoteStaticPlain
        override val handshakeHash: ByteArray = h.copyOf()
    }
}