package com.peachpasswords.sync

import kotlinx.coroutines.async
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.runBlocking
import org.bouncycastle.crypto.params.X25519PrivateKeyParameters
import org.junit.Assert.assertArrayEquals
import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.Test
import java.security.SecureRandom

class NoiseKKTest {
    private fun keys(seed: Int? = null): NoiseKKKeyPair {
        val privateKey = if (seed == null) {
            X25519PrivateKeyParameters(SecureRandom())
        } else {
            X25519PrivateKeyParameters(ByteArray(32) { seed.toByte() }, 0)
        }
        return NoiseKKKeyPair(privateKey.generatePublicKey().encoded, privateKey.encoded)
    }

    @Test
    fun `KK mutually authenticates and carries transport frames`() = runBlocking {
        val initiator = keys()
        val responder = keys()
        val aToB = Channel<ByteArray>(Channel.UNLIMITED)
        val bToA = Channel<ByteArray>(Channel.UNLIMITED)
        val initiatorHandshake = async {
            noiseKKInitiator(
                initiator.secretKey,
                initiator.publicKey,
                responder.publicKey,
                { aToB.send(it) },
                { bToA.receive() },
            )
        }
        val responderHandshake = async {
            noiseKKResponder(
                responder.secretKey,
                responder.publicKey,
                initiator.publicKey,
                { bToA.send(it) },
                { aToB.receive() },
            )
        }
        val a = initiatorHandshake.await()
        val b = responderHandshake.await()

        assertArrayEquals(a.handshakeHash, b.handshakeHash)
        assertEquals("one", String(b.decrypt(a.encrypt("one".encodeToByteArray()))))
        assertEquals("two", String(a.decrypt(b.encrypt("two".encodeToByteArray()))))
    }

    @Test
    fun `KK handshake hash matches browser cross-platform vector`() = runBlocking {
        val initiator = keys(0x01)
        val responder = keys(0x02)
        val initiatorEphemeral = keys(0x03)
        val responderEphemeral = keys(0x04)
        val aToB = Channel<ByteArray>(Channel.UNLIMITED)
        val bToA = Channel<ByteArray>(Channel.UNLIMITED)

        val initiatorHandshake = async {
            noiseKKInitiator(
                initiator.secretKey,
                initiator.publicKey,
                responder.publicKey,
                { aToB.send(it) },
                { bToA.receive() },
                initiatorEphemeral,
            )
        }
        val responderHandshake = async {
            noiseKKResponder(
                responder.secretKey,
                responder.publicKey,
                initiator.publicKey,
                { bToA.send(it) },
                { aToB.receive() },
                responderEphemeral,
            )
        }

        val a = initiatorHandshake.await()
        val b = responderHandshake.await()
        assertEquals(
            "dced1b411d7dc25716d8fc7fdc6b49418ae5648f323df5fb20a105f3c0acabab",
            pubkeyToHex(a.handshakeHash),
        )
        assertArrayEquals(a.handshakeHash, b.handshakeHash)
    }

    @Test
    fun `KK rejects an initiator key mismatch without fallback`() = runBlocking {
        val initiator = keys()
        val attacker = keys()
        val responder = keys()
        val aToB = Channel<ByteArray>(Channel.UNLIMITED)
        val bToA = Channel<ByteArray>(Channel.UNLIMITED)

        val initiatorHandshake = async {
            runCatching {
                noiseKKInitiator(
                    initiator.secretKey,
                    initiator.publicKey,
                    responder.publicKey,
                    { aToB.send(it) },
                    { bToA.receive() },
                )
            }
        }
        val responderHandshake = async {
            runCatching {
                noiseKKResponder(
                    responder.secretKey,
                    responder.publicKey,
                    attacker.publicKey,
                    { bToA.send(it) },
                    { aToB.receive() },
                )
            }
        }

        val failure = responderHandshake.await().exceptionOrNull()
        assertTrue(failure?.message?.contains("decrypt", ignoreCase = true) == true)
        initiatorHandshake.cancel()
    }

    // PB-26-014 Android parity: the nine-entry small-order X25519 blocklist
    // lives in the shared Rust core, so this test only proves the JNI driver
    // surfaces its rejection. The last two entries produce non-zero DH output,
    // so only the core's small-order guard (not the all-zero-output check)
    // rejects them.
    private val smallOrderRemoteStatics = listOf(
        "0000000000000000000000000000000000000000000000000000000000000000",
        "0100000000000000000000000000000000000000000000000000000000000000",
        "e0eb7a7c3b41b8ae1656e3faf19fc46ada098deb9c32b1fd866205165f49b880",
        "5f9c95bca3508c24b1d0b1559c83ef5b04445cc4581c8e86d8224eddd09f1157",
        "ecffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff7f",
        "edffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff7f",
        "eeffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff7f",
        "cdeb7a7c3b41b8ae1656e3faf19fc46ada098deb9c32b1fd866205165f49b880",
        "4c9c95bca3508c24b1d0b1559c83ef5b04445cc4581c8e86d8224eddd09f11d7",
    )

    @Test
    fun `KK initiator rejects every small-order remote static at the DH`() {
        val local = keys(0x0b)
        for (entry in smallOrderRemoteStatics) {
            val session = RustNoiseSession(
                RustNoiseNative.nativeCreateKkInitiator(
                    local.secretKey,
                    local.publicKey,
                    entry.hexToBytes(),
                ),
            )
            try {
                val failure = runCatching { session.writeHandshake() }.exceptionOrNull()
                assertTrue(
                    "small-order remote static $entry must fail the handshake at the DH",
                    failure != null,
                )
            } finally {
                session.close()
            }
        }
    }

    @Test
    fun `KK accepts an honest remote static`() {
        val local = keys(0x0b)
        val responder = keys(0x0c)
        val session = RustNoiseSession(
            RustNoiseNative.nativeCreateKkInitiator(
                local.secretKey,
                local.publicKey,
                responder.publicKey,
            ),
        )
        try {
            session.writeHandshake()
        } finally {
            session.close()
        }
    }
}

private fun String.hexToBytes(): ByteArray =
    ByteArray(length / 2) { index ->
        ((Character.digit(this[index * 2], 16) shl 4) or
            Character.digit(this[index * 2 + 1], 16)).toByte()
    }