package io.flamenet.e2e

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

/**
 * ML-KEM-768, the half of the hybrid handshake Android did not have.
 *
 * 🔴 **A session is post-quantum only when both ends are.** Every Android
 * device registered `proto: 1`, so iPhone's implementation protected no real
 * conversation — the property was true of the code and false of the product.
 * These tests are about the things that would let that keep being true while
 * looking fixed: a seed that round-trips to a different key, a size that is
 * wrong in a way nothing notices until a peer cannot decrypt.
 */
class PQKEMTest {

    @Test
    fun `a fresh keypair has the sizes FIPS 203 fixes`() {
        val pair = PQKEM.generate()
        assertEquals(PQKEM.SEED_LEN, pair.seed.size)
        assertEquals(PQKEM.PUBLIC_KEY_LEN, pair.publicKey.size)
    }

    @Test
    fun `encapsulate and decapsulate agree on a secret`() {
        val bob = PQKEM.generate()
        val out = PQKEM.encapsulate(bob.publicKey)

        assertEquals(PQKEM.CIPHERTEXT_LEN, out.ciphertext.size)
        assertEquals(PQKEM.SHARED_SECRET_LEN, out.sharedSecret.size)
        assertContentEquals(out.sharedSecret, PQKEM.decapsulate(out.ciphertext, bob.seed))
    }

    /**
     * 🔴 **The seed is the whole persistence story.** Only 64 bytes go into the
     * vault, and both the private key and the published public half are
     * re-derived from them on every load — so a seed that produced a different
     * public key on the second derivation would mean a device publishing a key
     * it cannot decapsulate with, and every peer who used it silently
     * unreachable.
     */
    @Test
    fun `a seed derives the same keypair every time`() {
        val first = PQKEM.generate()
        val again = PQKEM.fromSeed(first.seed)

        assertContentEquals(first.publicKey, again.publicKey)
        // And a secret encapsulated to one opens with the other.
        val out = PQKEM.encapsulate(first.publicKey)
        assertContentEquals(out.sharedSecret, PQKEM.decapsulate(out.ciphertext, again.seed))
    }

    /**
     * The seed path and the library's own generator must build the same key.
     *
     * ⚠️ Not a tautology. `fromSeed` reconstructs a private key from `d ‖ z` by
     * a different route than `generateKeyPair` takes, and a disagreement
     * between the two would show up only as a peer on another implementation
     * being unable to talk to us.
     */
    @Test
    fun `the seed constructor matches the library's own generator`() {
        val (seed, publicKey) = PQKEM.generateWithLibrary()
        assertEquals(PQKEM.SEED_LEN, seed.size, "the library stores a 64-byte FIPS 203 seed")
        assertContentEquals(publicKey, PQKEM.fromSeed(seed).publicKey)
    }

    @Test
    fun `two keypairs are not the same keypair`() {
        assertFalse(PQKEM.generate().publicKey.contentEquals(PQKEM.generate().publicKey))
    }

    /**
     * ⚠️ **Decapsulation does not fail on a tampered ciphertext**, by design:
     * ML-KEM's implicit rejection returns a well-formed but different secret
     * rather than an error. Asserted rather than assumed, because every caller
     * has to be written for it — the mismatch surfaces as the first message
     * failing to decrypt, which is why §12.5 also binds the ciphertext into the
     * session's associated data.
     */
    @Test
    fun `a tampered ciphertext yields a different secret, not an error`() {
        val bob = PQKEM.generate()
        val out = PQKEM.encapsulate(bob.publicKey)
        val tampered = out.ciphertext.copyOf().also { it[0] = (it[0].toInt() xor 1).toByte() }

        val got = PQKEM.decapsulate(tampered, bob.seed)
        assertEquals(PQKEM.SHARED_SECRET_LEN, got.size)
        assertFalse(got.contentEquals(out.sharedSecret))
    }

    @Test
    fun `the wrong seed opens nothing`() {
        val bob = PQKEM.generate()
        val out = PQKEM.encapsulate(bob.publicKey)
        assertFalse(PQKEM.decapsulate(out.ciphertext, PQKEM.generate().seed)
                        .contentEquals(out.sharedSecret))
    }

    /**
     * Sizes are checked at the door. A wrong length can only ever fail on every
     * peer, and it fails far away from here — so it is named where it happens
     * rather than discovered as a contact who stopped working.
     */
    @Test
    fun `wrong sizes are refused where they can still be explained`() {
        assertFailsWith<PQKEM.Malformed> { PQKEM.fromSeed(ByteArray(32)) }
        assertFailsWith<PQKEM.Malformed> { PQKEM.encapsulate(ByteArray(1183)) }
        assertFailsWith<PQKEM.Malformed> {
            PQKEM.decapsulate(ByteArray(1087), PQKEM.generate().seed)
        }
        assertFailsWith<PQKEM.Malformed> {
            PQKEM.decapsulate(ByteArray(PQKEM.CIPHERTEXT_LEN), ByteArray(63))
        }
    }

    /** Every encapsulation to one key is fresh. A repeat would be a broken RNG. */
    @Test
    fun `encapsulating twice to the same key gives two different ciphertexts`() {
        val bob = PQKEM.generate()
        val a = PQKEM.encapsulate(bob.publicKey)
        val b = PQKEM.encapsulate(bob.publicKey)
        assertFalse(a.ciphertext.contentEquals(b.ciphertext))
        assertFalse(a.sharedSecret.contentEquals(b.sharedSecret))
        // Both still open.
        assertContentEquals(a.sharedSecret, PQKEM.decapsulate(a.ciphertext, bob.seed))
        assertContentEquals(b.sharedSecret, PQKEM.decapsulate(b.ciphertext, bob.seed))
    }

    /** A secret with no entropy in it would pass every test above. */
    @Test
    fun `the shared secret is not all zeroes`() {
        val out = PQKEM.encapsulate(PQKEM.generate().publicKey)
        assertTrue(out.sharedSecret.any { it != 0.toByte() })
    }
}