package io.flamenet.e2e

import org.json.JSONObject
import java.io.File
import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertEquals
import kotlin.test.assertNotEquals

/**
 * The hybrid handshake against the published known-answer vector.
 *
 * 🔴 **This is the only test here worth much.** Every other assertion about
 * ML-KEM in this module compares Bouncy Castle to itself: seeds round-trip,
 * ciphertexts open, sizes are right — all true of an implementation that is
 * internally consistent and agrees with nobody. The vector was generated by the
 * normative JavaScript engine and is replayed by the Swift one, so matching it
 * is the claim that actually matters: a phone can start a post-quantum session
 * with an iPhone, and neither has to take our word for it.
 *
 * The specific thing at risk is the **seed**. FIPS 203 keygen takes `d ‖ z`,
 * and this engine stores only those 64 bytes and re-derives everything from
 * them — so if BC ordered or expanded them differently from CryptoKit, every
 * cross-platform session would fail at the first message with nothing pointing
 * at the cause. The vector settles it byte for byte.
 */
class PQVectorTest {

    private val vector: JSONObject by lazy {
        // Resolved rather than hardcoded, for the same reason PublishedVectorsTest
        // does it: this module lives both as a subtree of the app repo and as a
        // repo of its own beside the engine.
        //
        // ⚠️ Fails rather than skips when nothing is found. A conformance test
        // that quietly stops conforming is worth less than no test at all —
        // the file listing still says the handshake is checked.
        val candidates = listOfNotNull(
            System.getenv("FLAMENET_PQ_VECTORS"),
            "../../engine/test/vectors/pqx3dh-v2.json",
            "../../../flamenet-e2e/test/vectors/pqx3dh-v2.json",
            "../../flamenet-e2e/test/vectors/pqx3dh-v2.json",
        )
        val found = candidates.map(::File).firstOrNull { it.isFile }
            ?: error("no pqx3dh-v2.json found; looked in $candidates. " +
                     "Set FLAMENET_PQ_VECTORS to point at it.")
        JSONObject(found.readText())
    }

    private fun hex(bytes: ByteArray) = bytes.joinToString("") { "%02x".format(it) }

    /**
     * 🔴 **The cross-implementation check.** The vector publishes a seed and
     * the public key the normative engine derived from it; deriving a different
     * one here is the failure that would make every Android↔iPhone post-quantum
     * session fail on the first message, and nowhere near this code.
     */
    @Test
    fun `a published seed derives the published public key`() {
        val archive = vector.getJSONObject("bobArchive")
        for (which in listOf("pqSpk", "pqOtp")) {
            val entry = if (which == "pqSpk") archive.getJSONObject("pqSpk")
                        else archive.getJSONObject("pqOtp").getJSONObject("1")
            val seed = Bytes.unb64(entry.getString("seed"))
            val published = Bytes.unb64(entry.getString("pub"))

            assertEquals(PQKEM.SEED_LEN, seed.size, "$which seed is d‖z")
            assertContentEquals(published, PQKEM.fromSeed(seed).publicKey,
                                "$which: this engine derives a different key from the same seed")
        }
    }

    /**
     * The whole responder side, end to end, against the published answer.
     *
     * Everything is exercised at once and that is the point: the four
     * agreements, their order, the KEM secret going last, the v2 info string,
     * and the associated data. A vector that only checked the KEM would pass
     * while the concatenation order was wrong — and the two engines would still
     * never talk.
     */
    @Test
    fun `the responder derives the published session key and associated data`() {
        val archive = vector.getJSONObject("bobArchive")
        val env = vector.getJSONObject("env")
        val x3dh = env.getJSONObject("x3dh")
        val expect = vector.getJSONObject("expect")

        fun priv(path: JSONObject) = X25519.privateKeyFromPkcs8(Bytes.unb64(path.getString("priv")))

        val ciphertext = Bytes.unb64(env.getJSONObject("pq").getString("ct"))
        val pqSeed = Bytes.unb64(archive.getJSONObject("pqOtp").getJSONObject("1").getString("seed"))
        assertEquals("opk", env.getJSONObject("pq").getString("kind"),
                     "this vector uses the one-time PQ prekey")

        val sk = X3DH.respond(
            identityPrivateDH = priv(archive.getJSONObject("ikDH")),
            signedPreKeyPrivate = priv(archive.getJSONObject("spk")),
            oneTimePrivate = priv(archive.getJSONObject("otp").getJSONObject("1")),
            initiatorIdentityDH = Bytes.unb64(x3dh.getString("ik_dh")),
            initiatorEphemeral = Bytes.unb64(x3dh.getString("ek")),
            pq = pqSeed to ciphertext,
        )
        assertEquals(expect.getString("sk"), hex(sk))

        val ad = X3DH.associatedData(
            first = Bytes.unb64(x3dh.getString("ik_sig")),
            second = Bytes.unb64(archive.getJSONObject("ikSig").getString("pub")),
            pqCiphertext = ciphertext,
        )
        assertEquals(expect.getString("ad"), hex(ad))
    }

    /**
     * 🔴 **A downgrade must not produce a key the other side accepts.** v1's
     * `DH1‖DH2‖DH3‖DH4` and v2's `DH1‖DH2‖DH3‖SS` are both 128 bytes, so
     * without the info string changing with the version they are
     * indistinguishable KDF inputs (§12.5). Asserted against the published
     * answer rather than against another run of our own code.
     */
    @Test
    fun `running the same handshake without the KEM derives a different key`() {
        val archive = vector.getJSONObject("bobArchive")
        val x3dh = vector.getJSONObject("env").getJSONObject("x3dh")

        fun priv(path: JSONObject) = X25519.privateKeyFromPkcs8(Bytes.unb64(path.getString("priv")))

        val classical = X3DH.respond(
            identityPrivateDH = priv(archive.getJSONObject("ikDH")),
            signedPreKeyPrivate = priv(archive.getJSONObject("spk")),
            oneTimePrivate = priv(archive.getJSONObject("otp").getJSONObject("1")),
            initiatorIdentityDH = Bytes.unb64(x3dh.getString("ik_dh")),
            initiatorEphemeral = Bytes.unb64(x3dh.getString("ek")),
        )
        assertNotEquals(vector.getJSONObject("expect").getString("sk"), hex(classical))
    }

    /**
     * ⚠️ The associated data binds the ciphertext, so a relay that swapped one
     * cannot leave the session looking untouched. ML-KEM's implicit rejection
     * means decapsulation itself will not complain.
     */
    @Test
    fun `changing one byte of the ciphertext changes the associated data`() {
        val ciphertext = Bytes.unb64(vector.getJSONObject("env").getJSONObject("pq").getString("ct"))
        val tampered = ciphertext.copyOf().also { it[7] = (it[7].toInt() xor 1).toByte() }
        val a = X3DH.associatedData(ByteArray(32), ByteArray(32), ciphertext)
        val b = X3DH.associatedData(ByteArray(32), ByteArray(32), tampered)
        assertNotEquals(hex(a), hex(b))
    }

    /** A v1 session's AD is the two identity keys and nothing else. */
    @Test
    fun `without a ciphertext the associated data is just the two identities`() {
        val a = ByteArray(32) { 1 }
        val b = ByteArray(32) { 2 }
        assertContentEquals(a + b, X3DH.associatedData(a, b, null))
    }
}