package org.meshbay.client

import org.json.JSONObject
import org.junit.Assert.assertTrue
import org.junit.Test
import org.meshbay.client.keys.Kdf
import org.meshbay.client.keys.Keyring
import org.meshbay.client.keys.Transcripts
import java.io.File

/**
 * The keyring and the transcripts against meshbay-hub/tests/vectors/keyring.json,
 * which the desktop keyring writes and the specification reproduces
 * (test_keyring_vectors.py). Every deterministic field byte for byte, every
 * refusal with its message: a bundle this sealed is one the page and the
 * desktop open, and the reverse.
 */
class KeyringVectorsTest {
    private var passed = 0
    private val failures = ArrayList<String>()
    private fun check(label: String, ok: Boolean, detail: () -> String = { "" }) {
        if (ok) passed++ else failures.add("$label ${detail()}")
    }
    private fun <T> eq(label: String, got: T, want: T) = check(label, got == want) { "got=$got want=$want" }

    @Test fun `every vector is reproduced`() {

        val v = JSONObject(VECTORS.readText())
        val input = v.getJSONObject("input")
        val kdf = v.getJSONObject("kdf")
        val bundles = v.getJSONObject("bundles")
        val now = input.getLong("now").toDouble()
        val userId = input.getString("userId")
        val nodePk = input.getString("nodePk")
        val username = input.getString("username")
    
        // The store, as SecretStore would hold it.
        var store = JSONObject()
        var nonces: ArrayDeque<ByteArray> = ArrayDeque()
        fun ring() = Keyring(
            load = { JSONObject(store.toString()) },
            save = { store = JSONObject(it.toString()) },
            transcripts = Transcripts { now },
            random = { n -> nonces.removeFirstOrNull() ?: ByteArray(n).also { java.security.SecureRandom().nextBytes(it) } },
        )
        val ring = ring()
    
        // 1. KDF chain.
        val salt = Kdf.sha256("meshbay:bundle:v2:$username".toByteArray()).copyOfRange(0, 16)
        eq("salt", Kdf.toHex(salt), kdf.getString("salt_hex"))
        ring.deriveSession(input.getString("password"), username, userId,
                           input.getString("pepperB64"), input.getInt("pepperVersion"))
        val masters = store.getJSONObject("masters").getJSONObject(userId)
        eq("argon2id", Kdf.toHex(Kdf.unb64(masters.getString("legacy"))), kdf.getString("argon2_hex"))
        eq("M", Kdf.toHex(Kdf.unb64(masters.getString("m"))), kdf.getString("master_hex"))
        eq("pepper version", masters.getInt("v"), input.getInt("pepperVersion"))
        eq("fingerprint", ring.currentFingerprint(userId), kdf.getString("master_fingerprint"))
        eq("node key", Kdf.toHex(Kdf.hkdf(Kdf.unb64(masters.getString("m")), "meshbay:bundle:v3|node|$nodePk")),
           kdf.getString("node_key_hex"))
        eq("playlist key", ring.playlistKey(userId), kdf.getString("playlist_key_b64"))
    
        // 2. Identity: placed in the store as keyring.js would leave it.
        val ident = v.getJSONObject("identity")
        store.getJSONObject("identities").put(userId, JSONObject().put(nodePk,
            JSONObject().put("ed", ident.getString("ed_pkcs8_b64")).put("x", ident.getString("x_pkcs8_b64"))
                .put("sealedWith", JSONObject.NULL)))
        val pub = ring.identity(userId, nodePk)!!
        eq("pkEd", pub.pkEdB64, ident.getJSONObject("public").getString("pkEdB64"))
        eq("pkX", pub.pkXB64, ident.getJSONObject("public").getString("pkXB64"))
    
        // 3. Bundles: sealed byte for byte, and opened.
        val fixedNonce = Kdf.hex(input.getString("fixedNonceHex"))
        nonces.addLast(fixedNonce)
        val sealed = ring.sealBundle(userId, nodePk)
        eq("bundle sealed with a fixed nonce", sealed.bundle, bundles.getString("fixed_nonce_bundle_b64"))
        eq("bundle fingerprint", sealed.fingerprint, bundles.getString("fixed_nonce_fingerprint"))
        nonces.addLast(fixedNonce)
        eq("recovery bundle", ring.sealRecovery(userId, nodePk, input.getString("mnemonic"), username),
           bundles.getString("recovery_fixed_nonce_b64"))
    
        fun opensTo(label: String, block: (Keyring) -> Keyring.Pub) {
            val saved = store
            store = JSONObject().put("masters", JSONObject().put(userId, masters)).put("identities", JSONObject())
                .put("access", JSONObject())
            try {
                val p = block(ring())
                check(label, p.pkEdB64 == pub.pkEdB64 && p.pkXB64 == pub.pkXB64) { "opened to another identity" }
                check("$label leaves the identity unsealed", ring().identity(userId, nodePk)?.sealedWith == null)
            } catch (e: Exception) {
                check(label, false) { e.toString() }
            } finally { store = saved }
        }
        opensTo("desktop bundle (fixed nonce) opens") { it.openBundle(userId, nodePk, bundles.getString("fixed_nonce_bundle_b64")) }
        opensTo("MBK2 legacy bundle opens") { it.openBundle(userId, nodePk, bundles.getString("legacy_mbk2_b64")) }
        val bogus = Kdf.b64("MBK3".toByteArray() + ByteArray(30) { 1 })
        opensTo("recovery copy opens through the fallback") {
            it.openBundle(userId, nodePk, bogus, bundles.getString("recovery_fixed_nonce_b64"),
                          input.getString("mnemonic"), username)
        }
        // A bundle Kotlin seals (random nonce) must open — round trip.
        val ours = ring.sealBundle(userId, nodePk).bundle
        opensTo("a bundle sealed here opens here") { it.openBundle(userId, nodePk, ours) }
        // A retired format is refused, never tried against the recovery key.
        try {
            ring.openBundle(userId, nodePk, Kdf.b64("MBK1xxxxxxxxxxxxxxxxxxxxxxxxxxxxx".toByteArray()))
            check("retired format refused", false)
        } catch (e: Keyring.FormatRetired) { check("retired format refused", true) }
    
        // Browser access off: nothing sealed.
        ring.setBrowserAccess(userId, false)
        try { ring.sealBundle(userId, nodePk); check("no bundle while browser access is off", false) }
        catch (e: IllegalStateException) { check("no bundle while browser access is off", e.message!!.startsWith("Refused")) }
        ring.setBrowserAccess(userId, true)
    
        // 4. Agreement.
        val ag = v.getJSONObject("agreement")
        eq("X25519 agreement", ring.shared(userId, nodePk, ag.getString("peer_x_pub_b64")), ag.getString("shared_b64"))
    
        // 5. Transcripts and signatures.
        val tr = Transcripts { now }
        val ctx = Transcripts.Ctx(userId, nodePk, pub.pkEdB64, pub.pkXB64)
        val kinds = v.getJSONObject("transcripts")
        for (kind in kinds.keys()) {
            val t = kinds.getJSONObject(kind)
            eq("transcript $kind", Kdf.toHex(tr.forKind(kind, t.getJSONObject("fields"), ctx)), t.getString("transcript_hex"))
            eq("signature $kind", ring.signAs(userId, nodePk, kind, t.getJSONObject("fields")), t.getString("signature_b64"))
        }
        val refusals = v.getJSONArray("refusals")
        for (i in 0 until refusals.length()) {
            val r = refusals.getJSONObject(i)
            try {
                tr.forKind(r.getString("kind"), r.getJSONObject("fields"), ctx)
                check("refusal '${r.getString("label")}'", false) { "was signed" }
            } catch (e: Transcripts.Refused) {
                eq("refusal '${r.getString("label")}' message", e.message, r.getString("error"))
            }
        }
    
        // Sign-out drops M and keeps the identities.
        ring.forgetSession(userId)
        check("sign-out drops M", !ring.hasSession(userId))
        check("sign-out keeps the identity", ring.identity(userId, nodePk) != null)
        assertTrue("vector failures:\n" + failures.joinToString("\n"), failures.isEmpty())
        assertTrue("too few checks ran: $passed", passed >= 55)
    }

    companion object {
        // Unit tests run with the module directory as working directory.
        val VECTORS = File("../../meshbay-hub/tests/vectors/keyring.json")
    }
}