package org.meshbay.client.keys

import org.bouncycastle.crypto.agreement.X25519Agreement
import org.bouncycastle.crypto.params.Ed25519PrivateKeyParameters
import org.bouncycastle.crypto.params.X25519PrivateKeyParameters
import org.bouncycastle.crypto.params.X25519PublicKeyParameters
import org.bouncycastle.crypto.signers.Ed25519Signer
import org.json.JSONObject
import java.security.SecureRandom
import org.meshbay.client.keys.Kdf.b64
import org.meshbay.client.keys.Kdf.hkdf
import org.meshbay.client.keys.Kdf.unb64

/**
 * The account's keys on Android: the bundle master key `M` and the identity on
 * every node, held here and never handed to the page. A port of
 * meshbay-client/src/keyring.js — same state shape (masters, identities,
 * access), same formats, same refusals — so the two read side by side.
 *
 * Storage is injected (load/save of one JSON object), as on desktop; the app
 * keeps it in SecretStore (Keystore-wrapped). `random` is injected so the
 * vectors can pin a nonce.
 */
class Keyring(
    private val load: () -> JSONObject?,
    private val save: (JSONObject) -> Unit,
    private val transcripts: Transcripts = Transcripts(),
    private val random: (Int) -> ByteArray = { n -> ByteArray(n).also { SecureRandom().nextBytes(it) } },
) {
    class Pub(val pkEdB64: String, val pkXB64: String, val sealedWith: String? = null)
    class Sealed(val bundle: String, val fingerprint: String)
    class FormatRetired : IllegalStateException("bundle_format_retired")

    private class Master(val m: ByteArray, val v: Int)
    private class Identity(val ed: String, val x: String)

    // In memory: the key of a passphrase change not yet accepted by the hub.
    private val pending = HashMap<String, Master>()

    private fun state(): JSONObject {
        val s = load() ?: JSONObject()
        for (k in listOf("masters", "identities", "access")) if (!s.has(k)) s.put(k, JSONObject())
        return s
    }

    private fun master(userId: String, usePending: Boolean = false): Master {
        if (usePending) pending[userId]?.let { return it }
        val m = state().getJSONObject("masters").optJSONObject(userId)
            ?: throw IllegalStateException("no bundle key in this session")
        return Master(unb64(m.getString("m")), m.getInt("v"))
    }

    private fun fingerprint(m: ByteArray) = Kdf.toHex(Kdf.sha256(m)).substring(0, 16)

    private fun stored(userId: String, nodePk: String): Identity {
        val id = state().getJSONObject("identities").optJSONObject(userId)?.optJSONObject(nodePk)
            ?: throw IllegalStateException("no identity for this node")
        return Identity(id.getString("ed"), id.getString("x"))
    }

    private fun publicOf(id: Identity) = Pub(
        b64(Kdf.edFromPkcs8(unb64(id.ed)).generatePublicKey().encoded),
        b64(Kdf.xFromPkcs8(unb64(id.x)).generatePublicKey().encoded),
    )

    private fun accessOn(userId: String) = state().getJSONObject("access").opt(userId) != false
    private fun needAccess(userId: String) {
        if (!accessOn(userId)) throw IllegalStateException("Refused: browser access is off for this account")
    }

    private fun keep(userId: String, nodePk: String, put: (JSONObject) -> Unit) {
        val s = state()
        val ids = s.getJSONObject("identities")
        val forUser = ids.optJSONObject(userId) ?: JSONObject().also { ids.put(userId, it) }
        val entry = forUser.optJSONObject(nodePk) ?: JSONObject().also { forUser.put(nodePk, it) }
        put(entry)
        save(s)
    }

    private fun aad(userId: String, nodePk: String) =
        "meshbay:bundle:v3|$userId|$nodePk".toByteArray(Charsets.UTF_8)

    // Exactly JSON.stringify({ skEd, skX }): base64 needs no escaping, and the
    // sealed bytes are then comparable with the desktop's in tests.
    private fun plaintext(id: Identity) = "{\"skEd\":\"${id.ed}\",\"skX\":\"${id.x}\"}"

    private fun seal(id: Identity, key: ByteArray, userId: String, nodePk: String, pepperVersion: Int): String {
        val nonce = random(12)
        val ct = Kdf.gcmSeal(key, nonce, plaintext(id).toByteArray(Charsets.UTF_8), aad(userId, nodePk))
        return b64(MAGIC + byteArrayOf((pepperVersion and 0xff).toByte()) + nonce + ct)
    }

    private fun parse(plain: ByteArray): Identity {
        val o = JSONObject(String(plain, Charsets.UTF_8))
        return Identity(o.getString("skEd"), o.getString("skX"))
    }

    private fun open(bundleB64: String, key: ByteArray, userId: String, nodePk: String): Identity {
        val raw = unb64(bundleB64)
        if (raw.size < 4 || !raw.copyOfRange(0, 4).contentEquals(MAGIC)) throw FormatRetired()
        return parse(Kdf.gcmOpen(key, raw.copyOfRange(5, 17), raw.copyOfRange(17, raw.size), aad(userId, nodePk)))
    }

    /** TRANSITIONAL — MBK2: "MBK2" ‖ nonce ‖ AES-GCM under the Argon2 key, no AAD. */
    private fun openLegacy(bundleB64: String, key: ByteArray): Identity {
        val raw = unb64(bundleB64)
        return parse(Kdf.gcmOpen(key, raw.copyOfRange(4, 16), raw.copyOfRange(16, raw.size), null))
    }

    private fun fromMnemonic(mnemonic: String): ByteArray {
        val clean = mnemonic.replace(Regex("[^A-Za-z2-7]"), "").uppercase()
        var bits = 0
        var value = 0
        val out = ArrayList<Byte>()
        for (ch in clean) {
            value = (value shl 5) or B32.indexOf(ch)
            bits += 5
            if (bits >= 8) { out.add(((value ushr (bits - 8)) and 0xff).toByte()); bits -= 8 }
        }
        if (out.size < 32) throw IllegalArgumentException("recovery key too short")
        return out.subList(0, 32).toByteArray()
    }

    // ── The API, in keyring.js order ────────────────────────────────────────

    fun hasSession(userId: String) = state().getJSONObject("masters").has(userId)

    /** `M` from the passphrase and the pepper — one Argon2 run, as in the page. */
    fun deriveSession(password: String, username: String, userId: String, pepperB64: String?,
                      pepperVersion: Int?, pendingChange: Boolean = false): Boolean {
        if (userId.isEmpty() || pepperB64.isNullOrEmpty()) {
            throw IllegalStateException("the hub did not provide the bundle pepper")
        }
        val salt = Kdf.sha256("meshbay:bundle:v2:$username".toByteArray(Charsets.UTF_8)).copyOfRange(0, 16)
        val a = Kdf.argon2id(password, salt)
        val m = hkdf(a + unb64(pepperB64), "meshbay:bundle-master:v3|$userId")
        val v = if (pepperVersion == null || pepperVersion == 0) 1 else pepperVersion
        if (pendingChange) { pending[userId] = Master(m, v); return true }
        val s = state()
        // `legacy` (TRANSITIONAL): the Argon2 key MBK2 bundles were sealed under.
        s.getJSONObject("masters").put(userId, JSONObject().put("m", b64(m)).put("v", v).put("legacy", b64(a)))
        save(s)
        return true
    }

    fun commitPending(userId: String): Boolean {
        val p = pending[userId] ?: return false
        val s = state()
        val legacy = s.getJSONObject("masters").optJSONObject(userId)?.optString("legacy", "")
        val entry = JSONObject().put("m", b64(p.m)).put("v", p.v)
        if (!legacy.isNullOrEmpty()) entry.put("legacy", legacy)
        s.getJSONObject("masters").put(userId, entry)
        save(s)
        pending.remove(userId)
        return true
    }

    fun dropPending(userId: String) = pending.remove(userId) != null

    /** Sign-out: `M` goes. The identities stay — they are this device's. */
    fun forgetSession(userId: String): Boolean {
        val s = state()
        s.getJSONObject("masters").remove(userId)
        save(s)
        pending.remove(userId)
        return true
    }

    fun identity(userId: String, nodePk: String): Pub? {
        val id = state().getJSONObject("identities").optJSONObject(userId)?.optJSONObject(nodePk) ?: return null
        val pub = publicOf(Identity(id.getString("ed"), id.getString("x")))
        val sw = id.opt("sealedWith")
        return Pub(pub.pkEdB64, pub.pkXB64, if (sw is String) sw else null)
    }

    fun openBundle(userId: String, nodePk: String, bundleEnc: String, recoveryEnc: String? = null,
                   recoveryMnemonic: String? = null, username: String? = null): Pub {
        val raw = unb64(bundleEnc)
        if (raw.size >= 4 && raw.copyOfRange(0, 4).contentEquals(LEGACY_MAGIC)) {
            val legacy = state().getJSONObject("masters").optJSONObject(userId)?.optString("legacy", "")
            if (legacy.isNullOrEmpty()) throw IllegalStateException("no_legacy_key")
            val id = openLegacy(bundleEnc, unb64(legacy))
            keepIdentity(userId, nodePk, id)
            return publicOf(id)
        }
        val m = master(userId).m
        val id = try {
            open(bundleEnc, hkdf(m, "meshbay:bundle:v3|node|$nodePk"), userId, nodePk)
        } catch (e: Exception) {
            if (e is FormatRetired || recoveryEnc.isNullOrEmpty() || recoveryMnemonic.isNullOrEmpty()) throw e
            val rk = hkdf(fromMnemonic(recoveryMnemonic), "meshbay:recovery:v1:${username ?: ""}")
            open(recoveryEnc, rk, userId, nodePk)
        }
        keepIdentity(userId, nodePk, id)
        return publicOf(id)
    }

    private fun keepIdentity(userId: String, nodePk: String, id: Identity) =
        keep(userId, nodePk) { it.put("ed", id.ed).put("x", id.x).put("sealedWith", JSONObject.NULL) }

    fun mint(userId: String, nodePk: String): Pub {
        val rnd = SecureRandom()
        val id = Identity(b64(Kdf.edToPkcs8(Ed25519PrivateKeyParameters(rnd))),
                          b64(Kdf.xToPkcs8(X25519PrivateKeyParameters(rnd))))
        keepIdentity(userId, nodePk, id)
        return publicOf(id)
    }

    /** The identity sealed for its node, under `M` (or the pending one). */
    fun sealBundle(userId: String, nodePk: String, usePending: Boolean = false): Sealed {
        needAccess(userId)
        val m = master(userId, usePending)
        val bundle = seal(stored(userId, nodePk), hkdf(m.m, "meshbay:bundle:v3|node|$nodePk"), userId, nodePk, m.v)
        return Sealed(bundle, fingerprint(m.m))
    }

    /** The recovery copy: sealed under the recovery key, owing nothing to `M`. */
    fun sealRecovery(userId: String, nodePk: String, mnemonic: String, username: String): String {
        needAccess(userId)
        val rk = hkdf(fromMnemonic(mnemonic), "meshbay:recovery:v1:$username")
        return seal(stored(userId, nodePk), rk, userId, nodePk, 0)
    }

    fun markSealed(userId: String, nodePk: String, fp: String?): Boolean {
        keep(userId, nodePk) { it.put("sealedWith", if (fp.isNullOrEmpty()) JSONObject.NULL else fp) }
        return true
    }

    fun currentFingerprint(userId: String) = fingerprint(master(userId).m)

    /** Sign what `kind` names, built from `fields` (Transcripts). */
    fun signAs(userId: String, nodePk: String, kind: String, fields: JSONObject?): String {
        val id = stored(userId, nodePk)
        val pub = publicOf(id)
        val transcript = transcripts.forKind(kind, fields, Transcripts.Ctx(userId, nodePk, pub.pkEdB64, pub.pkXB64))
        val signer = Ed25519Signer()
        signer.init(true, Kdf.edFromPkcs8(unb64(id.ed)))
        signer.update(transcript, 0, transcript.size)
        return b64(signer.generateSignature())
    }

    fun shared(userId: String, nodePk: String, peerPkB64: String): String {
        val agreement = X25519Agreement()
        agreement.init(Kdf.xFromPkcs8(unb64(stored(userId, nodePk).x)))
        val out = ByteArray(32)
        agreement.calculateAgreement(X25519PublicKeyParameters(unb64(peerPkB64), 0), out, 0)
        return b64(out)
    }

    fun playlistKey(userId: String) = b64(hkdf(master(userId).m, "meshbay:playlists:v2"))

    fun browserAccess(userId: String) = accessOn(userId)
    fun setBrowserAccess(userId: String, on: Boolean): Boolean {
        val s = state()
        s.getJSONObject("access").put(userId, on)
        save(s)
        return on
    }

    companion object {
        private val MAGIC = "MBK3".toByteArray()
        // TRANSITIONAL — the format before MBK3, read once to be replaced.
        private val LEGACY_MAGIC = "MBK2".toByteArray()
        private const val B32 = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"
    }
}