package no.zyon.shade

import no.zyon.shade.approval.canonicalApprovalSigningBytes
import no.zyon.shade.backup.deriveBackupKey
import no.zyon.shade.blob.aeadOpen
import no.zyon.shade.blob.blobAadForSlot
import no.zyon.shade.blob.deriveBlobKey
import no.zyon.shade.blob.deriveBlobSigningSeed
import no.zyon.shade.blob.deriveBlobSlotId
import no.zyon.shade.blob.ed25519PublicKeyFromSeed
import no.zyon.shade.crypto.TinkProvider
import no.zyon.shade.fingerprint.computeFingerprint
import no.zyon.shade.group.encodeSenderHeader
import no.zyon.shade.group.senderKeyStep
import no.zyon.shade.group.senderSignedBytes
import no.zyon.shade.protocol.deriveInitialRootKey
import no.zyon.shade.protocol.kdfChainKey
import no.zyon.shade.protocol.kdfRootKey
import no.zyon.shade.serialization.StreamChunkWire
import no.zyon.shade.serialization.StreamChunkWireFormat
import no.zyon.shade.serialization.WireFormat
import no.zyon.shade.streams.aesGcmDecryptWithNonce
import no.zyon.shade.streams.aesGcmEncryptWithNonce
import no.zyon.shade.streams.buildChunkAad
import no.zyon.shade.streams.buildChunkNonce
import no.zyon.shade.streams.deriveLaneKey
import no.zyon.shade.streams.deriveStreamKey
import no.zyon.shade.types.PreKeyMessage
import no.zyon.shade.types.RatchetMessage
import no.zyon.shade.types.ShadeEnvelope
import org.json.JSONArray
import org.json.JSONObject
import org.junit.Assert.assertEquals
import org.junit.Assert.assertTrue
import org.junit.Test
import java.io.File
import java.nio.ByteBuffer
import javax.crypto.Cipher
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.SecretKeySpec

/**
 * Cross-platform test vectors. MUST match the TypeScript implementation
 * byte-for-byte, otherwise cross-platform messaging breaks.
 *
 * The test-vectors/ directory is at the root of the Shade monorepo.
 * Generated by scripts/generate-vectors.ts from the TypeScript implementation.
 */
class CrossPlatformVectorTest {

    private val crypto = TinkProvider()
    private val vectorsDir = File("../../test-vectors")
    private val expectedVersion = 2

    private fun fromHex(str: String): ByteArray {
        val bytes = ByteArray(str.length / 2)
        for (i in bytes.indices) {
            bytes[i] = ((Character.digit(str[i * 2], 16) shl 4) +
                Character.digit(str[i * 2 + 1], 16)).toByte()
        }
        return bytes
    }

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

    private data class VectorFile(val version: Int, val vectors: JSONArray)

    private fun loadVectors(name: String): JSONArray {
        val file = File(vectorsDir, name)
        val content = file.readText()
        val obj = JSONObject(content)
        val version = obj.getInt("version")
        assertEquals("Unexpected vector schema version in $name", expectedVersion, version)
        return obj.getJSONArray("vectors")
    }

    private fun encodeRatchetHeader(
        dhPublicKey: ByteArray,
        previousCounter: Int,
        counter: Int,
    ): ByteArray {
        val buf = ByteBuffer.allocate(40)
        buf.put(dhPublicKey)
        buf.putInt(previousCounter)
        buf.putInt(counter)
        return buf.array()
    }

    private fun aesGcmEncryptDeterministic(
        key: ByteArray,
        nonce: ByteArray,
        plaintext: ByteArray,
        aad: ByteArray,
    ): ByteArray {
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        val spec = GCMParameterSpec(128, nonce)
        cipher.init(Cipher.ENCRYPT_MODE, SecretKeySpec(key, "AES"), spec)
        cipher.updateAAD(aad)
        return cipher.doFinal(plaintext)
    }

    @Test
    fun hkdfVectorsMatch() {
        val vectors = loadVectors("hkdf.json")
        for (i in 0 until vectors.length()) {
            val v = vectors.getJSONObject(i)
            val out = crypto.hkdf(
                fromHex(v.getString("ikm")),
                fromHex(v.getString("salt")),
                v.getString("info").toByteArray(Charsets.UTF_8),
                v.getInt("length"),
            )
            assertEquals(v.getString("output"), hex(out))
        }
    }

    @Test
    fun kdfChainVectorsMatch() {
        val vectors = loadVectors("kdf-chain.json")

        val rootVec = vectors.getJSONObject(0)
        val rootResult = kdfRootKey(
            crypto,
            fromHex(rootVec.getString("rootKey")),
            fromHex(rootVec.getString("dhOutput")),
        )
        assertEquals(rootVec.getString("newRootKey"), hex(rootResult.newRootKey))
        assertEquals(rootVec.getString("chainKey"), hex(rootResult.chainKey))

        val chainVec = vectors.getJSONObject(1)
        val chainResult = kdfChainKey(crypto, fromHex(chainVec.getString("chainKey")))
        assertEquals(chainVec.getString("newChainKey"), hex(chainResult.newChainKey))
        assertEquals(chainVec.getString("messageKey"), hex(chainResult.messageKey))
    }

    @Test
    fun x3dhVectorsMatch() {
        val vectors = loadVectors("x3dh.json")
        for (i in 0 until vectors.length()) {
            val v = vectors.getJSONObject(i)
            val secretsArray = v.getJSONArray("secrets")
            val secrets = (0 until secretsArray.length()).map { fromHex(secretsArray.getString(it)) }
            val rootKey = deriveInitialRootKey(crypto, secrets)
            assertEquals(v.getString("rootKey"), hex(rootKey))
        }
    }

    @Test
    fun fingerprintVectorsMatch() {
        val vectors = loadVectors("fingerprint.json")
        for (i in 0 until vectors.length()) {
            val v = vectors.getJSONObject(i)
            val fp = computeFingerprint(
                crypto,
                fromHex(v.getString("signingKey")),
                fromHex(v.getString("dhKey")),
            )
            assertEquals(v.getString("fingerprint"), fp)
        }
    }

    @Test
    fun wireFormatRatchetVectorsMatch() {
        val vectors = loadVectors("wire-format.json")
        var found = false
        for (i in 0 until vectors.length()) {
            val v = vectors.getJSONObject(i)
            if (v.optString("kind") != "ratchet") continue
            found = true
            val m = v.getJSONObject("message")

            val msg = RatchetMessage(
                dhPublicKey = fromHex(m.getString("dhPublicKey")),
                previousCounter = m.getInt("previousCounter"),
                counter = m.getInt("counter"),
                ciphertext = fromHex(m.getString("ciphertext")),
                nonce = fromHex(m.getString("nonce")),
            )
            val envelope = ShadeEnvelope(
                type = ShadeEnvelope.EnvelopeType.RATCHET,
                content = msg,
                timestamp = 0,
                senderAddress = "",
            )
            val encoded = WireFormat.encodeEnvelope(envelope)
            assertEquals(v.getString("encoded"), hex(encoded))

            val decoded = WireFormat.decodeEnvelope(encoded)
            assertEquals(ShadeEnvelope.EnvelopeType.RATCHET, decoded.type)
            val rm = decoded.content as RatchetMessage
            assertEquals(msg.counter, rm.counter)
            assertEquals(hex(msg.ciphertext), hex(rm.ciphertext))
        }
        assertTrue("No ratchet wire vectors found", found)
    }

    @Test
    fun wireFormatPreKeyVectorsMatch() {
        val vectors = loadVectors("wire-format.json")
        var matched = 0
        for (i in 0 until vectors.length()) {
            val v = vectors.getJSONObject(i)
            if (v.optString("kind") != "prekey") continue
            matched++
            val m = v.getJSONObject("message")
            val inner = m.getJSONObject("inner")

            val innerMsg = RatchetMessage(
                dhPublicKey = fromHex(inner.getString("dhPublicKey")),
                previousCounter = inner.getInt("previousCounter"),
                counter = inner.getInt("counter"),
                ciphertext = fromHex(inner.getString("ciphertext")),
                nonce = fromHex(inner.getString("nonce")),
            )
            val preKeyId: Int? = if (m.isNull("preKeyId")) null else m.getInt("preKeyId")
            val pre = PreKeyMessage(
                registrationId = m.getInt("registrationId"),
                preKeyId = preKeyId,
                signedPreKeyId = m.getInt("signedPreKeyId"),
                ephemeralKey = fromHex(m.getString("ephemeralKey")),
                identityDHKey = fromHex(m.getString("identityDHKey")),
                message = innerMsg,
            )
            val envelope = ShadeEnvelope(
                type = ShadeEnvelope.EnvelopeType.PREKEY,
                content = pre,
                timestamp = 0,
                senderAddress = "",
            )
            val encoded = WireFormat.encodeEnvelope(envelope)
            assertEquals(v.getString("encoded"), hex(encoded))

            val decoded = WireFormat.decodeEnvelope(encoded)
            assertEquals(ShadeEnvelope.EnvelopeType.PREKEY, decoded.type)
            val dm = decoded.content as PreKeyMessage
            assertEquals(pre.registrationId, dm.registrationId)
            assertEquals(pre.preKeyId, dm.preKeyId)
            assertEquals(pre.signedPreKeyId, dm.signedPreKeyId)
            assertEquals(hex(pre.ephemeralKey), hex(dm.ephemeralKey))
            assertEquals(hex(innerMsg.ciphertext), hex(dm.message.ciphertext))
        }
        assertTrue("Expected at least 2 prekey vectors", matched >= 2)
    }

    private fun findVector(arr: JSONArray, prefix: String): JSONObject {
        for (i in 0 until arr.length()) {
            val o = arr.getJSONObject(i)
            if (o.getString("description").startsWith(prefix)) return o
        }
        throw AssertionError("Vector with description prefix '$prefix' not found")
    }

    @Test
    fun streamsVectorsMatch() {
        val vectors = loadVectors("streams.json")

        // 1. deriveStreamKey
        val sk = findVector(vectors, "deriveStreamKey")
        val streamSecret = fromHex(sk.getString("streamSecret"))
        val streamId = fromHex(sk.getString("streamId"))
        val streamKey = deriveStreamKey(crypto, streamSecret, streamId)
        assertEquals(sk.getString("streamKey"), hex(streamKey))

        // 2. deriveLaneKey
        val lk = findVector(vectors, "deriveLaneKey")
        val lkStreamKey = fromHex(lk.getString("streamKey"))
        val lkStreamId = fromHex(lk.getString("streamId"))
        val lanes = lk.getJSONArray("lanes")
        for (i in 0 until lanes.length()) {
            val lane = lanes.getJSONObject(i)
            val laneId = lane.getLong("laneId")
            val k = deriveLaneKey(crypto, lkStreamKey, lkStreamId, laneId)
            assertEquals(lane.getString("laneKey"), hex(k))
        }

        // 3. buildChunkNonce
        val nv = findVector(vectors, "buildChunkNonce")
        val nonces = nv.getJSONArray("nonces")
        for (i in 0 until nonces.length()) {
            val n = nonces.getJSONObject(i)
            val laneId = n.getLong("laneId")
            val seq = java.lang.Long.parseUnsignedLong(n.getString("seq"))
            val out = buildChunkNonce(laneId, seq)
            assertEquals(n.getString("nonce"), hex(out))
        }

        // 4. buildChunkAad
        val av = findVector(vectors, "buildChunkAad")
        val avStreamId = fromHex(av.getString("streamId"))
        val cases = av.getJSONArray("cases")
        for (i in 0 until cases.length()) {
            val c = cases.getJSONObject(i)
            val laneId = c.getLong("laneId")
            val seq = java.lang.Long.parseUnsignedLong(c.getString("seq"))
            val isLast = c.getBoolean("isLast")
            val out = buildChunkAad(avStreamId, laneId, seq, isLast)
            assertEquals(c.getString("aad"), hex(out))
        }

        // 5. End-to-end chunk encrypt + decrypt
        val ev = findVector(vectors, "End-to-end chunk encrypt")
        val laneKey = fromHex(ev.getString("laneKey"))
        val nonce = fromHex(ev.getString("nonce"))
        val aad = fromHex(ev.getString("aad"))
        val plaintext = fromHex(ev.getString("plaintext"))
        val ct = aesGcmEncryptWithNonce(laneKey, nonce, plaintext, aad)
        assertEquals(ev.getString("ciphertext"), hex(ct))
        val pt = aesGcmDecryptWithNonce(laneKey, nonce, fromHex(ev.getString("ciphertext")), aad)
        assertEquals(ev.getString("plaintext"), hex(pt))

        // 6. Wire 0x11 envelope encode/decode
        val wv = findVector(vectors, "Wire 0x11")
        val wire = StreamChunkWire(
            streamId = fromHex(wv.getString("streamId")),
            laneId = wv.getLong("laneId"),
            seq = java.lang.Long.parseUnsignedLong(wv.getString("seq")),
            isLast = wv.getBoolean("isLast"),
            nonce = fromHex(wv.getString("nonce")),
            aad = fromHex(wv.getString("extraAad")),
            ciphertext = fromHex(wv.getString("ciphertext")),
        )
        val encoded = StreamChunkWireFormat.encodeStreamChunk(wire)
        assertEquals(wv.getString("encoded"), hex(encoded))

        val decoded = StreamChunkWireFormat.decodeStreamChunk(encoded)
        assertEquals(hex(wire.streamId), hex(decoded.streamId))
        assertEquals(wire.laneId, decoded.laneId)
        assertEquals(wire.seq, decoded.seq)
        assertEquals(wire.isLast, decoded.isLast)
        assertEquals(hex(wire.nonce), hex(decoded.nonce))
        assertEquals(hex(wire.ciphertext), hex(decoded.ciphertext))

        // 7. Envelope-type inspector
        assertEquals(
            StreamChunkWireFormat.EnvelopeKind.STREAM_CHUNK,
            StreamChunkWireFormat.inspectEnvelopeType(encoded),
        )
    }

    @Test
    fun backupVectorsMatch() {
        val vectors = loadVectors("backup.json")

        val kv = findVector(vectors, "Backup v1: HKDF")
        val backupKey = deriveBackupKey(crypto, kv.getString("passphrase"), fromHex(kv.getString("salt")))
        assertEquals(kv.getString("backupKey"), hex(backupKey))

        val ev = findVector(vectors, "Backup v1: AES-256-GCM")
        val ct = aesGcmEncryptDeterministic(
            fromHex(ev.getString("backupKey")),
            fromHex(ev.getString("nonce")),
            fromHex(ev.getString("plaintext")),
            ByteArray(0),
        )
        assertEquals(ev.getString("ciphertext"), hex(ct))

        val pt = crypto.aesGcmDecrypt(
            fromHex(ev.getString("backupKey")),
            fromHex(ev.getString("ciphertext")),
            fromHex(ev.getString("nonce")),
            null,
        )
        assertEquals(ev.getString("plaintext"), hex(pt))
    }

    @Test
    fun groupSenderKeyVectorsMatch() {
        val vectors = loadVectors("group.json")

        // 1. Header AAD
        val hv = findVector(vectors, "Sender header AAD")
        val aad = encodeSenderHeader(
            hv.getString("groupId"),
            hv.getString("senderAddress"),
            hv.getInt("iteration"),
        )
        assertEquals(hv.getString("aad"), hex(aad))

        // 2. Sender-key step
        val sv = findVector(vectors, "Sender-key step")
        val step = senderKeyStep(
            crypto,
            fromHex(sv.getString("chainKey")),
            sv.getString("groupId"),
            sv.getString("senderAddress"),
            sv.getInt("iteration"),
        )
        assertEquals(sv.getString("newChainKey"), hex(step.newChainKey))
        assertEquals(sv.getString("messageKey"), hex(step.messageKey))
        assertEquals(sv.getString("aad"), hex(step.aad))

        val ct = aesGcmEncryptDeterministic(
            step.messageKey,
            fromHex(sv.getString("nonce")),
            fromHex(sv.getString("plaintext")),
            step.aad,
        )
        assertEquals(sv.getString("ciphertext"), hex(ct))

        // 3. Ed25519 verify on the recorded signature
        val signed = senderSignedBytes(step.aad, ct)
        val ok = crypto.verify(
            fromHex(sv.getString("signingPublicKey")),
            signed,
            fromHex(sv.getString("signature")),
        )
        assertTrue("Sender-key signature verification failed", ok)

        // 4. Decrypt roundtrip
        val pt = crypto.aesGcmDecrypt(
            step.messageKey,
            fromHex(sv.getString("ciphertext")),
            fromHex(sv.getString("nonce")),
            step.aad,
        )
        assertEquals(sv.getString("plaintext"), hex(pt))
    }

    @Test
    fun storageHkdfVectorsMatch() {
        val vectors = loadVectors("storage-hkdf.json")

        val sv = findVector(vectors, "Storage HKDF: storageKey")
        val storageKey = crypto.hkdf(
            fromHex(sv.getString("masterKey")),
            ByteArray(0),
            "shade-storage-v1".toByteArray(Charsets.UTF_8),
            32,
        )
        assertEquals(sv.getString("storageKey"), hex(storageKey))

        val fv = findVector(vectors, "Storage HKDF: fieldKey")
        val fStorageKey = fromHex(fv.getString("storageKey"))
        val fields = fv.getJSONArray("fields")
        for (i in 0 until fields.length()) {
            val f = fields.getJSONObject(i)
            val info = "shade-field-v1:${f.getString("table")}:${f.getString("column")}"
                .toByteArray(Charsets.UTF_8)
            val k = crypto.hkdf(fStorageKey, ByteArray(0), info, 32)
            assertEquals(f.getString("fieldKey"), hex(k))
        }

        val nv = findVector(vectors, "Storage HKDF: rowNonce")
        val rowKey = fromHex(nv.getString("rowKey"))
        val nonces = nv.getJSONArray("nonces")
        for (i in 0 until nonces.length()) {
            val n = nonces.getJSONObject(i)
            val info = "shade-row-nonce-v1:${n.getString("table")}:${n.getString("pk")}"
                .toByteArray(Charsets.UTF_8)
            val out = crypto.hkdf(rowKey, ByteArray(0), info, 12)
            assertEquals(n.getString("nonce"), hex(out))
        }
    }

    @Test
    fun blobKdfAndAeadVectorsMatch() {
        val vectors = loadVectors("blob.json")
        var kdfMatched = 0
        var aeadMatched = 0
        for (i in 0 until vectors.length()) {
            val v = vectors.getJSONObject(i)
            val desc = v.getString("description")
            if (desc.startsWith("V4.9 blob KDF")) {
                kdfMatched++
                val masterKey = fromHex(v.getString("masterKey"))
                val app = v.getString("app")
                val slotId = deriveBlobSlotId(crypto, masterKey, app)
                assertEquals(v.getString("slotId"), hex(slotId))
                assertEquals(v.getString("blobKey"), hex(deriveBlobKey(crypto, masterKey, app)))
                val seed = deriveBlobSigningSeed(crypto, masterKey, app)
                assertEquals(v.getString("signingSeed"), hex(seed))
                assertEquals(v.getString("ownerPubkey"), hex(ed25519PublicKeyFromSeed(seed)))
            } else if (desc.startsWith("V4.9 blob AEAD")) {
                aeadMatched++
                val key = fromHex(v.getString("key"))
                val slotIdHex = v.getString("slotIdHex")
                val expectedPlaintext = fromHex(v.getString("plaintext"))
                val wire = fromHex(v.getString("wire"))
                val aad = blobAadForSlot(slotIdHex)
                val opened = aeadOpen(key, wire, aad)
                assertEquals(hex(expectedPlaintext), hex(opened))
            }
        }
        assertTrue("KDF vectors expected", kdfMatched >= 3)
        assertTrue("AEAD vectors expected", aeadMatched >= 2)
    }

    @Test
    fun approvalSigningPayloadVectorsMatch() {
        val vectors = loadVectors("approval.json")
        var payloadMatched = 0
        var e2eMatched = 0
        for (i in 0 until vectors.length()) {
            val v = vectors.getJSONObject(i)
            val desc = v.getString("description")
            if (desc.startsWith("V4.10 approval signing payload")) {
                payloadMatched++
                val out = canonicalApprovalSigningBytes(
                    domain = v.getString("domain"),
                    requestId = v.getString("requestId"),
                    hostFingerprint = v.getString("hostFingerprint"),
                    requestingDeviceFingerprint = v.getString("requestingDeviceFingerprint"),
                    decision = v.getString("decision"),
                )
                assertEquals(v.getString("signingPayload"), hex(out))
            } else if (desc.startsWith("V4.10 approval Ed25519 sign/verify")) {
                e2eMatched++
                val seed = fromHex(v.getString("seed"))
                val pubkey = fromHex(v.getString("publicKey"))
                assertEquals(hex(pubkey), hex(ed25519PublicKeyFromSeed(seed)))

                val req = v.getJSONObject("request")
                val payload = canonicalApprovalSigningBytes(
                    domain = req.getString("domain"),
                    requestId = req.getString("requestId"),
                    hostFingerprint = req.getString("hostFingerprint"),
                    requestingDeviceFingerprint = req.getString("requestingDeviceFingerprint"),
                    decision = req.getString("decision"),
                )
                assertEquals(v.getString("signingPayload"), hex(payload))

                // Verify the TS-generated signature against our pubkey + payload.
                // This is the load-bearing parity check: a Kotlin-implemented
                // verifyProxyApproval running against a TS-signed approval
                // succeeds.
                val sig = fromHex(v.getString("signature"))
                val ok = crypto.verify(pubkey, payload, sig)
                assertTrue("Ed25519 verify of TS-signed approval failed", ok)

                // And: Kotlin signs the same payload with the same seed and
                // produces a sig the TS pubkey verifies. Ed25519 is
                // deterministic, so the sig bytes also match exactly.
                val mySig = crypto.sign(seed, payload)
                assertEquals(v.getString("signature"), hex(mySig))
            }
        }
        assertTrue("payload vectors expected", payloadMatched >= 3)
        assertTrue("e2e sign/verify vector expected", e2eMatched >= 1)
    }

    @Test
    fun ratchetStepRoundtripMatches() {
        val vectors = loadVectors("ratchet-step.json")
        assertTrue("ratchet-step vectors expected", vectors.length() > 0)

        for (i in 0 until vectors.length()) {
            val v = vectors.getJSONObject(i)
            val inputs = v.getJSONObject("inputs")
            val derived = v.getJSONObject("derived")

            val rootKey = fromHex(inputs.getString("rootKey"))
            val dhSendPriv = fromHex(inputs.getString("dhSendPrivateKey"))
            val dhSendPub = fromHex(inputs.getString("dhSendPublicKey"))
            val dhRemotePub = fromHex(inputs.getString("dhRemotePublicKey"))
            val plaintext = fromHex(inputs.getString("plaintext"))
            val nonce = fromHex(inputs.getString("nonce"))
            val previousCounter = inputs.getInt("previousCounter")
            val counter = inputs.getInt("counter")

            // 1. DH
            val dhOutput = crypto.x25519(dhSendPriv, dhRemotePub)
            assertEquals(derived.getString("dhOutput"), hex(dhOutput))

            // 2. kdfRootKey
            val root = kdfRootKey(crypto, rootKey, dhOutput)
            assertEquals(derived.getString("newRootKey"), hex(root.newRootKey))
            assertEquals(derived.getString("chainKey"), hex(root.chainKey))

            // 3. kdfChainKey
            val chain = kdfChainKey(crypto, root.chainKey)
            assertEquals(derived.getString("newChainKey"), hex(chain.newChainKey))
            assertEquals(derived.getString("messageKey"), hex(chain.messageKey))

            // 4. Header AAD
            val aad = encodeRatchetHeader(dhSendPub, previousCounter, counter)
            assertEquals(derived.getString("aad"), hex(aad))

            // 5. AES-GCM encrypt with fixed nonce
            val ciphertext = aesGcmEncryptDeterministic(chain.messageKey, nonce, plaintext, aad)
            assertEquals(v.getString("ciphertext"), hex(ciphertext))

            // 6. Roundtrip decrypt
            val recovered = crypto.aesGcmDecrypt(
                chain.messageKey,
                fromHex(v.getString("ciphertext")),
                nonce,
                aad,
            )
            assertEquals(inputs.getString("plaintext"), hex(recovered))
        }
    }
}
