package io.flamenet.e2e

import org.json.JSONObject
import kotlin.test.Test
import kotlin.test.assertContentEquals
import kotlin.test.assertEquals
import kotlin.test.assertFailsWith
import kotlin.test.assertNotNull
import kotlin.test.assertNull
import kotlin.test.assertTrue

/**
 * Two devices negotiating the hybrid handshake, end to end.
 *
 * 🔴 **A session is post-quantum only when both ends are.** The primitive
 * working proves nothing on its own: what matters is that a real registration
 * publishes usable material, that a real peer picks it up, that the ciphertext
 * survives the envelope, and that the responder finds the seed again. Any one
 * of those failing leaves a device that advertises `proto: 2` and derives a
 * different key from everybody — which surfaces as messages that will not
 * decrypt and points nowhere near the cause.
 *
 * The other half of these tests is about what happens when the relay lies,
 * because everything in a bundle passes through one.
 */
class HybridSessionTest {

    /** A device that has published post-quantum material, as a peer sees it. */
    private fun published(device: DeviceIdentity, withOneTime: Boolean = true): PeerBundle {
        val opk = device.generateOneTimePreKeys(1).first()
        val wire = device.registrationWire()
        val pqOtp = if (withOneTime) device.generatePQOneTimePreKeys(1).first() else null
        val bundle = JSONObject()
            .put("device_id", wire.getString("device_id"))
            .put("ik_dh", wire.getString("ik_dh"))
            .put("ik_sig", wire.getString("ik_sig"))
            .put("spk_id", wire.getInt("spk_id"))
            .put("spk", wire.getString("spk"))
            .put("spk_sig", wire.getString("spk_sig"))
            .put("proto", wire.getInt("proto"))
            .put("caps_sig", wire.getString("caps_sig"))
            .put("pqspk_id", wire.getInt("pqspk_id"))
            .put("pqspk", wire.getString("pqspk"))
            .put("pqspk_sig", wire.getString("pqspk_sig"))
            .put("opk", JSONObject().put("id", opk.id).put("pub", Bytes.b64(opk.publicKey)))
        if (pqOtp != null) {
            bundle.put("pqopk", JSONObject()
                .put("id", pqOtp.id)
                .put("pub", Bytes.b64(pqOtp.publicKey))
                .put("sig", Bytes.b64(pqOtp.signature)))
        }
        return PeerBundle.fromWire(bundle)
    }

    /**
     * A peer that has proven v2 and only v2 — the released fleet.
     *
     * ⚠️ `copy(proto = 2)` alone would NOT do: the capability signature commits
     * to the version, so a bundle whose proto was edited reads as v1 to every
     * peer. It has to be re-signed for 2 to be a genuine v2 device.
     */
    private fun publishedV2(device: DeviceIdentity, withOneTime: Boolean = true): PeerBundle =
        published(device, withOneTime).copy(
            proto = 2, capabilitySignature = device.capabilitySignature(2))

    /** A v1 device: exactly what this engine published before ML-KEM existed. */
    private fun publishedV1(device: DeviceIdentity): PeerBundle {
        val opk = device.generateOneTimePreKeys(1).first()
        val wire = device.registrationWire(postQuantum = false)
        return PeerBundle.fromWire(JSONObject()
            .put("device_id", wire.getString("device_id"))
            .put("ik_dh", wire.getString("ik_dh"))
            .put("ik_sig", wire.getString("ik_sig"))
            .put("spk_id", wire.getInt("spk_id"))
            .put("spk", wire.getString("spk"))
            .put("spk_sig", wire.getString("spk_sig"))
            .put("proto", wire.getInt("proto"))
            .put("opk", JSONObject().put("id", opk.id).put("pub", Bytes.b64(opk.publicKey))))
    }

    // MARK: - The happy path

    @Test
    fun `two post-quantum devices hold a post-quantum conversation`() {
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        val session = E2ESession.initiate(alice, published(bob))
        assertEquals(3, session.version,
                     "both ends implement the ratchet, so this must be v3")

        val envelope = session.encrypt("post-quantum hello")
        assertNotNull(envelope.pq, "the KEM ciphertext has to travel with the handshake")
        assertEquals("opk", envelope.pq!!.kind, "a fresh pool means the one-time key")
        assertEquals(PQKEM.CIPHERTEXT_LEN, envelope.pq!!.ciphertext.size)

        // Through the wire, because that is where a missing field hides.
        val delivered = Envelope.fromPayload(envelope.toPayload())
        val bobSide = E2ESession.respond(bob, delivered)

        // 🔴 **3, and asserting 2 here is what hid a real bug for weeks.**
        // The initiator is v3 two assertions above; a responder that answers
        // the same handshake as v2 disagrees with its own peer about which
        // protocol the session runs. Kotlin-to-Kotlin never noticed, because
        // both ends drive the ratchet off their own PQ state and ignore the
        // label — so this test round-tripped happily while encoding the defect.
        //
        // A v3 peer that does enforce it does notice: iOS refuses a ratchet
        // step with no post-quantum material as a downgrade, so every message
        // Android sent into an iOS-initiated session was dropped on arrival and
        // the conversation ran one way only.
        assertEquals(3, bobSide.version,
                     "answering a v3 handshake makes a v3 session, not a v2 one")
        assertEquals("post-quantum hello", bobSide.decrypt(delivered))
        assertContentEquals(session.associatedData, bobSide.associatedData,
                            "both sides must bind the same ciphertext digest")

        // And the reply comes back on the same session.
        val reply = bobSide.encrypt("and hello back")
        assertEquals("and hello back", session.decrypt(Envelope.fromPayload(reply.toPayload())))
    }

    /**
     * ⚠️ **The last-resort key is a fallback, not a choice.** Every session
     * then shares one KEM key, which is a real weakening — but it is far better
     * than dropping to v1 silently, which is what a device with no last-resort
     * key would force.
     */
    @Test
    fun `a drained one-time pool still gives a post-quantum session`() {
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        val session = E2ESession.initiate(alice, published(bob, withOneTime = false))
        assertEquals(3, session.version)
        val envelope = session.encrypt("still hybrid")
        assertEquals("spk", envelope.pq!!.kind)
        assertEquals("still hybrid", E2ESession.respond(bob, envelope).decrypt(envelope))
    }

    // MARK: - Mixed fleets

    /** Somebody on an older build is still reachable, on v1. */
    @Test
    fun `a v2 device talking to a v1 device falls back cleanly`() {
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()

        val session = E2ESession.initiate(alice, publishedV1(bob))
        assertEquals(1, session.version)
        val envelope = session.encrypt("classical hello")
        assertNull(envelope.pq)
        assertEquals("classical hello", E2ESession.respond(bob, envelope).decrypt(envelope))
    }

    /** And the other direction: a v1 initiator reaching a v2 device. */
    @Test
    fun `a v1 initiator can still start a session with a v2 device`() {
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()
        // Alice's engine has ML-KEM but the bundle she was served does not
        // offer any, which is what a v1 peer looks like.
        val bundleWithoutPQ = publishedV1(bob)
        val session = E2ESession.initiate(alice, bundleWithoutPQ)
        val envelope = session.encrypt("hello")
        assertEquals(1, E2ESession.respond(bob, envelope).version)
    }

    // MARK: - What the relay is not allowed to do

    /**
     * 🔴 **The substitution this whole layer exists to prevent.** A relay that
     * can swap in an ML-KEM key of its own reads the post-quantum half of every
     * new session — exactly what `spk_sig` stops classically. A key whose
     * signature does not check out is treated as absent, so the session drops
     * to v1 rather than to a key the relay chose.
     */
    @Test
    fun `a substituted post-quantum prekey is ignored, not used`() {
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()
        val honest = published(bob)
        val relaysOwnKey = PQKEM.generate()

        val tampered = honest.copy(
            pqOneTimePreKey = relaysOwnKey.publicKey,
            pqSignedPreKey = relaysOwnKey.publicKey,
        )
        assertNull(tampered.provenPostQuantumPreKey(),
                   "an unsigned KEM key must not be encapsulated to")
        assertEquals(1, E2ESession.initiate(alice, tampered).version,
                     "and the session falls back to v1, not to the relay's key")
    }

    /**
     * ⚠️ A relay can also lie the other way — advertise a v2 device as v1 and
     * quietly strip the post-quantum half off every new session. `caps_sig` is
     * what makes the claim checkable; where the signature is missing the peer
     * reads as v1, which is the safe reading of "we cannot tell".
     */
    @Test
    fun `a bundle with no capability signature is treated as v1`() {
        val bob = DeviceIdentity.create()
        assertNull(published(bob).copy(capabilitySignature = null).provenPostQuantumPreKey())
    }

    /** A capability signature for a different version does not verify. */
    @Test
    fun `a capability signature does not transfer between versions`() {
        val bob = DeviceIdentity.create()
        // 🔴 A version this engine does not claim, so the edit cannot be
        // accidentally valid. This read `proto = 3` while the engine published
        // 2; the constant moving is what makes a hardcoded neighbour like that
        // stop testing anything.
        val forAnotherVersion = published(bob).copy(proto = DeviceIdentity.PROTOCOL_VERSION + 1)
        assertNull(forAnotherVersion.provenPostQuantumPreKey(),
                   "the signature commits to the version it claims")
    }

    /**
     * ⚠️ The one-time id is signed **big-endian**, and the id is inside the
     * signature so a relay cannot renumber a key to make a peer spend the wrong
     * one.
     */
    @Test
    fun `a renumbered one-time prekey does not verify`() {
        val bob = DeviceIdentity.create()
        val moved = published(bob).copy(pqOneTimePreKeyId = 99)
        // It falls back to the last-resort key rather than accepting the
        // renumbered one, which is exactly the intended degradation.
        assertEquals("spk", moved.provenPostQuantumPreKey()?.kind)
    }

    /**
     * 🔴 **The advertised version and the capability signature must agree.**
     *
     * A peer verifies `Ed25519(ik_sig, "FlamenetE2E_Caps_v2" ‖ uint8(proto))`
     * against the proto in the same bundle. Bumping the constant without the
     * signature following it produces a bundle every peer silently reads as v1
     * — `provenPostQuantumPreKey()` returns null and the whole post-quantum
     * path disappears, with nothing anywhere reporting an error.
     */
    @Test
    fun `the published bundle claims v3 and proves it`() {
        val device = DeviceIdentity.create()
        val wire = device.registrationWire()
        assertEquals(3, wire.getInt("proto"), "this engine implements the v3 ratchet")

        // Round-trip through the peer's own view, which is the only reading
        // that matters: it re-verifies the signature over the claimed version.
        val seen = published(device)
        assertEquals(3, seen.proto)
        assertNotNull(seen.provenPostQuantumPreKey(),
                      "the capability signature does not cover the version claimed")
    }

    /**
     * 🔴 **A v2 peer must not be sent a v3 ratchet step.**
     *
     * The released fleet advertises `proto: 2`: it does post-quantum key
     * AGREEMENT and has no post-quantum ratchet. Attaching `pqdh`/`pqct` to a
     * message for one of those devices breaks the conversation twice over —
     * their root step derives without `ss` and under the `_v1` info string
     * while ours used `_v3`, and their canonical header omits the two fields
     * entirely so the AAD does not match either. The first message never opens
     * and never will.
     *
     * ⚠️ Offering a PQ prekey is NOT the same claim as implementing §13. Every
     * v2 device offers one. The version has to be what gates this, which is
     * what §13.4 says and what the iOS engine already does.
     */
    @Test
    fun `a peer that only claims v2 is not sent a post-quantum ratchet step`() {
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()
        // A real v2 bundle: PQ prekeys published, capability signed for 2.
        val session = E2ESession.initiate(alice, publishedV2(bob))
        val envelope = session.encrypt("hello")
        assertNull(envelope.header.pqdh,
                   "a v2 peer was sent a v3 ratchet step it cannot read")
        assertNull(envelope.header.pqct,
                   "a v2 peer was sent a PQ ciphertext it cannot decapsulate")
    }

    // MARK: - Persistence

    /**
     * 🔴 The seed is the only thing stored, so a device that could not find it
     * again would publish a key it cannot decapsulate with — and every peer who
     * used that key would be silently unreachable.
     */
    @Test
    fun `post-quantum material survives an archive round trip`() {
        val bob = DeviceIdentity.create()
        val bundle = published(bob)
        val restored = DeviceIdentity.fromArchive(JSONObject(bob.archive().toString()))

        val session = E2ESession.initiate(DeviceIdentity.create(), bundle)
        val envelope = session.encrypt("after a restart")
        assertEquals("after a restart", E2ESession.respond(restored, envelope).decrypt(envelope))
    }

    /**
     * ⚠️ A device archived before this engine had ML-KEM has no such block, and
     * every existing install is one. It must load as a v1 device rather than
     * fail, and publish post-quantum material the next time it registers.
     */
    @Test
    fun `an archive from before post-quantum still loads`() {
        val old = DeviceIdentity.create()
        val archive = old.archive()
        archive.remove("pqSpk"); archive.remove("pqOtp"); archive.remove("nextPQId")

        val loaded = DeviceIdentity.fromArchive(JSONObject(archive.toString()))
        assertTrue(!loaded.hasPQMaterial)
        // Registering makes one, so the upgrade needs nothing from the user.
        assertTrue(loaded.registrationWire().getString("pqspk").isNotEmpty())
        assertTrue(loaded.hasPQMaterial)
    }

    /**
     * ⚠️ **A prekey we already spent is not an error.** A duplicate delivery of
     * a v2 prekey envelope names a one-time key this device has consumed, and
     * throwing there would turn a retransmission into a dead conversation. It
     * simply cannot be v2 — and then fails to decrypt, which the caller's three
     * chains already handle.
     */
    @Test
    fun `a replayed prekey envelope does not throw`() {
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()
        val envelope = E2ESession.initiate(alice, published(bob)).encrypt("once")

        assertEquals("once", E2ESession.respond(bob, envelope).decrypt(envelope))
        // The one-time key is gone now. The second attempt must fail as a
        // decrypt, not as an exception out of the handshake.
        val second = E2ESession.respond(bob, envelope)
        assertEquals(1, second.version, "no PQ seed left, so this cannot be v2")
        assertFailsWith<Exception> { second.decrypt(envelope) }
    }

    /** The one-time key is spent exactly once. That is what makes it one-time. */
    @Test
    fun `a one-time post-quantum prekey is consumed`() {
        val bob = DeviceIdentity.create()
        bob.generatePQOneTimePreKeys(2)
        assertEquals(2, bob.pqOneTimePreKeyCount)
        assertNotNull(bob.takePQPreKey("opk", 1))
        assertEquals(1, bob.pqOneTimePreKeyCount)
        assertNull(bob.takePQPreKey("opk", 1))
        // The last-resort key is not consumed, ever.
        bob.pqSignedPreKey()
        assertNotNull(bob.takePQPreKey("spk", 1))
        assertNotNull(bob.takePQPreKey("spk", 1))
    }

    // MARK: - The wire

    /**
     * ⚠️ **The v2 block is all-or-nothing**, and the relay refuses a partial
     * one. A device advertising `proto: 2` without a usable last-resort prekey
     * would be picked for v2 by every peer and then fail every handshake.
     */
    @Test
    fun `a post-quantum registration publishes the whole block`() {
        val wire = DeviceIdentity.create().registrationWire(pqOneTime = emptyList())
        assertEquals(DeviceIdentity.PROTOCOL_VERSION, wire.getInt("proto"))
        for (field in listOf("caps_sig", "pqspk", "pqspk_sig")) {
            assertTrue(wire.getString(field).isNotEmpty(), "$field must be published")
        }
        assertEquals(PQKEM.PUBLIC_KEY_LEN, Bytes.unb64(wire.getString("pqspk")).size)
    }

    /** And a caller that asks for v1 gets exactly the old body back. */
    @Test
    fun `a v1 registration publishes no post-quantum fields`() {
        val wire = DeviceIdentity.create().registrationWire(postQuantum = false)
        assertEquals(1, wire.getInt("proto"))
        for (field in listOf("caps_sig", "pqspk", "pqspk_sig", "pq_prekeys")) {
            assertTrue(wire.isNull(field) || !wire.has(field), "$field must be absent")
        }
    }

    /**
     * The KEM ciphertext rides the prekey envelope and nothing after it. Every
     * later message is byte-identical to v1, because the secret is already in
     * the root key and the associated data binds the ciphertext for the life of
     * the session (§12.6) — republishing 1088 bytes per message would be pure
     * cost.
     */
    @Test
    fun `the ciphertext travels once, not on every message`() {
        val alice = DeviceIdentity.create()
        val bob = DeviceIdentity.create()
        val session = E2ESession.initiate(alice, published(bob))

        val first = session.encrypt("one")
        assertNotNull(first.pq)
        // Still pending: the handshake repeats until the peer answers, because
        // any single message may be the first to arrive.
        assertNotNull(session.encrypt("two").pq)

        val bobSide = E2ESession.respond(bob, first)
        // ⚠️ Read it before replying. A responder that answers without
        // decrypting has no receiving chain to step from — the same mistake a
        // caller would make, so the test makes it once and not again.
        bobSide.decrypt(first)
        session.decrypt(bobSide.encrypt("answered"))
        assertNull(session.encrypt("three").pq, "answered, so the handshake stops repeating")
    }
}