package io.flamenet.e2e

import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertFalse
import kotlin.test.assertTrue

/**
 * X3DH, both directions.
 *
 * The property that matters is agreement: Alice computes a secret from her
 * private keys and Bob's public bundle, Bob computes one from his private keys
 * and what Alice sent, and the two must be identical. A wrong DH order, a
 * swapped argument, or a zero-filled DH4 all still produce 32 plausible bytes —
 * they just produce *different* 32 bytes on each side, which is why this checks
 * equality rather than shape.
 */
class X3DHTest {

    private fun raw(kp: java.security.KeyPair) = X25519.rawFromPublicKey(kp.public)

    @Test
    fun `both sides derive the same secret, with a one-time prekey`() {
        val aliceIdentity = X25519.generate()
        val bobIdentity = X25519.generate()
        val bobSigned = X25519.generate()
        val bobOneTime = X25519.generate()

        val alice = X3DH.initiate(
            identityPrivateDH = aliceIdentity.private,
            bundle = X3DH.Bundle(
                identityKeyDH = raw(bobIdentity),
                signedPreKey = raw(bobSigned),
                oneTimePreKey = raw(bobOneTime),
            ),
        )
        val bob = X3DH.respond(
            identityPrivateDH = bobIdentity.private,
            signedPreKeyPrivate = bobSigned.private,
            oneTimePrivate = bobOneTime.private,
            initiatorIdentityDH = raw(aliceIdentity),
            initiatorEphemeral = alice.ephemeralPublic,
        )
        assertContentEquals(alice.sharedSecret, bob, "the two sides disagree")
        assertTrue(alice.sharedSecret.size == 32)
    }

    @Test
    fun `both sides agree when the one-time prekey pool is drained`() {
        // A degraded but valid state: initial forward secrecy is lost, and the
        // session must still work. Omitting DH4 on one side while the other
        // zero-fills it is the classic way to break exactly this case.
        val aliceIdentity = X25519.generate()
        val bobIdentity = X25519.generate()
        val bobSigned = X25519.generate()

        val alice = X3DH.initiate(
            identityPrivateDH = aliceIdentity.private,
            bundle = X3DH.Bundle(raw(bobIdentity), raw(bobSigned), null),
        )
        val bob = X3DH.respond(
            identityPrivateDH = bobIdentity.private,
            signedPreKeyPrivate = bobSigned.private,
            oneTimePrivate = null,
            initiatorIdentityDH = raw(aliceIdentity),
            initiatorEphemeral = alice.ephemeralPublic,
        )
        assertContentEquals(alice.sharedSecret, bob)
    }

    @Test
    fun `using a one-time prekey changes the result`() {
        // Otherwise DH4 is contributing nothing and its absence would be
        // undetectable — which would make a drained pool silently equivalent to
        // a healthy one.
        val aliceIdentity = X25519.generate()
        val bobIdentity = X25519.generate()
        val bobSigned = X25519.generate()
        val bobOneTime = X25519.generate()
        val ephemeral = X25519.generate()

        val withOpk = X3DH.initiate(
            aliceIdentity.private,
            X3DH.Bundle(raw(bobIdentity), raw(bobSigned), raw(bobOneTime)),
            ephemeral,
        )
        val withoutOpk = X3DH.initiate(
            aliceIdentity.private,
            X3DH.Bundle(raw(bobIdentity), raw(bobSigned), null),
            ephemeral,
        )
        assertFalse(withOpk.sharedSecret.contentEquals(withoutOpk.sharedSecret))
    }

    @Test
    fun `a different responder derives a different secret`() {
        val aliceIdentity = X25519.generate()
        val bobIdentity = X25519.generate()
        val bobSigned = X25519.generate()
        val malloryIdentity = X25519.generate()
        val mallorySigned = X25519.generate()

        val toBob = X3DH.initiate(
            aliceIdentity.private, X3DH.Bundle(raw(bobIdentity), raw(bobSigned), null))
        val malloryTries = X3DH.respond(
            identityPrivateDH = malloryIdentity.private,
            signedPreKeyPrivate = mallorySigned.private,
            oneTimePrivate = null,
            initiatorIdentityDH = raw(aliceIdentity),
            initiatorEphemeral = toBob.ephemeralPublic,
        )
        assertFalse(toBob.sharedSecret.contentEquals(malloryTries))
    }
}