package io.flamenet.e2e

import org.json.JSONObject
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.SecretKey
import javax.crypto.spec.GCMParameterSpec

/**
 * Everything the app must remember, encrypted at rest.
 *
 * ⚠️ **This deliberately does not use [Aead].** `Aead` takes raw key bytes, and
 * the key here comes from the Android Keystore, which never hands out its key
 * material — that is the entire reason to use it. A design that needed the raw
 * bytes would quietly have to keep them somewhere, which is the thing being
 * avoided. So this works with a `SecretKey` handle and the platform `Cipher`.
 *
 * ⚠️ **Every record's name is bound into the AAD.** Without that, a blob is
 * portable between slots: someone with file access could paste the session for
 * one contact into another contact's slot, and the app would decrypt it
 * happily and attribute messages to the wrong person. The name is not a
 * filename here, it is part of the ciphertext's meaning.
 *
 * ⚠️ **Deleting has to actually delete.** "Deleted means deleted" is a promise
 * the product makes, so [forget] removes the record rather than marking it, and
 * [forgetEverything] takes the identity with it.
 */
class Vault(private val storage: Storage, private val key: SecretKey) {

    /** Wherever bytes live. A directory on a phone, a map in a test. */
    interface Storage {
        fun read(name: String): ByteArray?
        fun write(name: String, bytes: ByteArray)
        fun delete(name: String)
        fun names(): List<String>
    }

    class Corrupt(name: String, cause: Throwable?) :
        Exception("the vault record '$name' did not open", cause)

    companion object {
        private const val NONCE_LEN = 12
        private const val TAG_BITS = 128

        private const val IDENTITY = "identity"
        private const val WATERMARK = "watermark"
        private const val CREDENTIAL = "credential"
        private const val SESSION_PREFIX = "session/"
        private const val HISTORY_PREFIX = "history/"
        private const val TRUST_PREFIX = "trust/"
        private const val KEYLOG_PIN = "keylog-pin"
        private const val DELIVERY_KEY = "delivery-key"

        /** The two kinds of thread, named once so they cannot drift apart. */
        fun threadForUser(userId: Long) = "user:$userId"
        fun threadForGroup(groupId: String) = "group:$groupId"

        private val random = SecureRandom()
    }

    // -- Bytes ---------------------------------------------------------------

    private fun seal(name: String, plaintext: ByteArray): ByteArray {
        val nonce = ByteArray(NONCE_LEN).also(random::nextBytes)
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(Cipher.ENCRYPT_MODE, key, GCMParameterSpec(TAG_BITS, nonce))
        cipher.updateAAD(name.toByteArray(Charsets.UTF_8))
        return nonce + cipher.doFinal(plaintext)
    }

    private fun open(name: String, box: ByteArray): ByteArray {
        if (box.size <= NONCE_LEN) throw Corrupt(name, null)
        return try {
            val cipher = Cipher.getInstance("AES/GCM/NoPadding")
            cipher.init(
                Cipher.DECRYPT_MODE, key,
                GCMParameterSpec(TAG_BITS, box.copyOfRange(0, NONCE_LEN)),
            )
            cipher.updateAAD(name.toByteArray(Charsets.UTF_8))
            cipher.doFinal(box, NONCE_LEN, box.size - NONCE_LEN)
        } catch (e: Exception) {
            // Deliberately one exception for every cause. Which of "wrong key",
            // "wrong slot" and "flipped bit" it was is information an attacker
            // would like and the user cannot act on.
            throw Corrupt(name, e)
        }
    }

    private fun readJson(name: String): JSONObject? =
        storage.read(name)?.let { JSONObject(String(open(name, it), Charsets.UTF_8)) }

    private fun writeJson(name: String, json: JSONObject) =
        storage.write(name, seal(name, json.toString().toByteArray(Charsets.UTF_8)))

    // -- Identity ------------------------------------------------------------

    /**
     * ⚠️ Returning null here means "this is a new install", and the caller will
     * mint a fresh device — which every contact sees as "security keys
     * changed", the warning that is supposed to mean an attack. So a record
     * that exists but will not open must throw rather than read as absent.
     */
    fun loadIdentity(): DeviceIdentity? = readJson(IDENTITY)?.let(DeviceIdentity::fromArchive)

    fun saveIdentity(identity: DeviceIdentity) = writeJson(IDENTITY, identity.archive())

    fun hasIdentity(): Boolean = storage.read(IDENTITY) != null

    // -- The account credential ----------------------------------------------

    /**
     * Who we are signed in as, so a restart does not mean signing in again.
     *
     * ⚠️ **This is the long-lived credential, not the password.** The relay
     * mints it, returns it once and keeps only its hash, so this is the only
     * copy — losing it means signing in again, and leaking it means someone
     * else is signed in as this account until it is revoked. It lives here
     * rather than in shared preferences for exactly that reason.
     *
     * ⚠️ The relay address is stored with it and must be. The same credential
     * against a different relay is meaningless, and a client that remembered
     * the credential but not where it came from would present it to whatever
     * server happened to be configured next.
     */
    data class Credential(
        val relayBaseUrl: String,
        val username: String,
        val appPassword: String,
        val userId: Long,
    )

    fun loadCredential(): Credential? = readJson(CREDENTIAL)?.let {
        Credential(
            relayBaseUrl = it.getString("relay"),
            username = it.getString("username"),
            appPassword = it.getString("appPassword"),
            userId = it.getLong("userId"),
        )
    }

    fun saveCredential(credential: Credential) = writeJson(CREDENTIAL, JSONObject()
        .put("relay", credential.relayBaseUrl)
        .put("username", credential.username)
        .put("appPassword", credential.appPassword)
        .put("userId", credential.userId))

    // -- Sessions ------------------------------------------------------------

    private fun sessionName(peerDeviceId: String) = SESSION_PREFIX + peerDeviceId

    fun loadSession(peerDeviceId: String): E2ESession? =
        readJson(sessionName(peerDeviceId))?.let(E2ESession::fromArchive)

    fun saveSession(peerDeviceId: String, session: E2ESession) =
        writeJson(sessionName(peerDeviceId), session.archive())

    fun knownPeers(): List<String> =
        storage.names().filter { it.startsWith(SESSION_PREFIX) }.map { it.removePrefix(SESSION_PREFIX) }

    // -- Identity trust (§19) ------------------------------------------------
    //
    // ⚠️ **Encrypted at rest like everything else here, and that is the point
    // rather than a nicety.** These records are what the client compares a
    // served key against. Somewhere an attacker can rewrite them is somewhere
    // an attacker can pin their own key and have the comparison agree with
    // them for ever afterwards — silently, because a match is what "fine"
    // looks like.

    private fun trustName(userId: Long) = TRUST_PREFIX + userId

    /**
     * What we have seen this person's devices use, and whether a human checked.
     *
     * `identities` is every identity key pinned for them, base64. Plural
     * because a person is not a device: someone with a phone and a tablet has
     * two, and both are legitimately theirs.
     */
    data class Trust(
        val identities: List<String>,
        val verified: Boolean,
        /**
         * The global address this contact was first seen under, if any.
         *
         * ⚠️ **Pinned, not looked up per use.** The key log files a leaf under
         * whichever subject the account had at publication time, so verifying a
         * served key means recomputing against the address we already accepted
         * for this person — not against whatever the relay says today. Asking
         * again each time would let a relay change the answer and have every
         * check agree with the new one.
         *
         * Null means "no address seen", which is both "never claimed one" and
         * "the relay would not say" — the relay answers those identically on
         * purpose, so the client cannot tell them apart either.
         */
        val fnid: String? = null,
    )

    fun loadTrust(userId: Long): Trust? = readJson(trustName(userId))?.let { o ->
        val a = o.optJSONArray("identities") ?: org.json.JSONArray()
        Trust(
            identities = (0 until a.length()).map { a.getString(it) },
            verified = o.optBoolean("verified", false),
            fnid = o.optString("fnid").ifEmpty { null },
        )
    }

    fun saveTrust(userId: Long, trust: Trust) = writeJson(trustName(userId), JSONObject()
        .put("identities", org.json.JSONArray(trust.identities))
        .put("verified", trust.verified)
        .put("fnid", trust.fnid ?: JSONObject.NULL))

    fun forgetTrust(userId: Long) = storage.delete(trustName(userId))

    // -- Sealed sender (§14.4) -----------------------------------------------

    /**
     * This account's delivery key, and whether the relay has its digest yet.
     *
     * 🔴 **The key itself never goes to the relay — only `SHA-256(key)` does.**
     * It reaches contacts inside E2E messages, so the relay never sees it from
     * its owner, and that is the only reason an unauthenticated sealed
     * submission can be authorised without naming who is submitting. Sending
     * the raw key to the server would make the whole mechanism decorative.
     *
     * `published` tracks the digest upload separately from the key, because a
     * key generated and never published is one every contact would try to seal
     * to and every seal would be rejected.
     */
    data class Delivery(
        val key: ByteArray,
        val published: Boolean,
        /** Peer user id -> their delivery key, so we can reach them sealed. */
        val peers: Map<Long, ByteArray> = emptyMap(),
        /** Peers we have already handed OUR key to, so we send it once. */
        val toldPeers: Set<Long> = emptySet(),
    ) {
        override fun equals(other: Any?): Boolean =
            other is Delivery && key.contentEquals(other.key) && published == other.published &&
                peers.keys == other.peers.keys && toldPeers == other.toldPeers

        override fun hashCode(): Int = key.contentHashCode() * 31 + published.hashCode()
    }

    fun loadDelivery(): Delivery? = readJson(DELIVERY_KEY)?.let { o ->
        val peers = o.optJSONObject("peers") ?: JSONObject()
        val told = o.optJSONArray("told") ?: org.json.JSONArray()
        Delivery(
            key = java.util.Base64.getDecoder().decode(o.getString("key")),
            published = o.optBoolean("published", false),
            peers = peers.keys().asSequence().associate { k ->
                k.toLong() to java.util.Base64.getDecoder().decode(peers.getString(k))
            },
            toldPeers = (0 until told.length()).map { told.getLong(it) }.toSet(),
        )
    }

    fun saveDelivery(delivery: Delivery) {
        val peers = JSONObject()
        delivery.peers.forEach { (id, k) ->
            peers.put(id.toString(), java.util.Base64.getEncoder().encodeToString(k))
        }
        writeJson(DELIVERY_KEY, JSONObject()
            .put("key", java.util.Base64.getEncoder().encodeToString(delivery.key))
            .put("published", delivery.published)
            .put("peers", peers)
            .put("told", org.json.JSONArray(delivery.toldPeers.toList())))
    }

    // -- Key-log pin (§19.3) -------------------------------------------------

    /**
     * The head this device last accepted, and the reason a rewritten log is
     * detectable at all.
     *
     * Trust on first use: absent, the first head seen is taken and pinned.
     * Present, every later head must be provably an extension of it — so a
     * relay that rewrites history has to produce a consistency proof between
     * two trees where none exists, which is the one thing it cannot do.
     */
    data class KeyLogPin(val size: Int, val root: String)

    fun loadKeyLogPin(): KeyLogPin? = readJson(KEYLOG_PIN)?.let {
        KeyLogPin(it.getInt("size"), it.getString("root"))
    }

    fun saveKeyLogPin(pin: KeyLogPin) = writeJson(KEYLOG_PIN, JSONObject()
        .put("size", pin.size)
        .put("root", pin.root))

    // -- Inbox watermark -----------------------------------------------------

    /**
     * The highest inbox id already **decrypted and stored**.
     *
     * ⚠️ Advance it on receipt rather than after storing and a crash mid-batch
     * loses the rest of the batch silently — the relay will never offer those
     * ids again.
     */
    var inboxWatermark: Long
        get() = readJson(WATERMARK)?.optLong("id", 0) ?: 0
        set(value) = writeJson(WATERMARK, JSONObject().put("id", value))

    // -- Message history -----------------------------------------------------

    /**
     * Messages, encrypted at rest, keyed by the person rather than the device.
     *
     * ⚠️ **Keyed by thread, not by device.** A conversation is with someone —
     * or with a group — and their phone and tablet are the same conversation.
     * Filing by device shows one person as several threads that each hold half
     * the story. Thread ids are `user:<id>` and `group:<gid>`; see
     * [threadForUser] and [threadForGroup].
     *
     * ⚠️ The whole thread is rewritten on every append. That is fine for a
     * chat and would not be for a log; if a thread ever grows large enough for
     * this to matter, the fix is an append-only record format, not a bigger
     * buffer.
     */
    fun history(threadId: String): org.json.JSONArray =
        readJson(historyName(threadId))?.optJSONArray("messages") ?: org.json.JSONArray()

    fun append(threadId: String, message: JSONObject) {
        val existing = history(threadId)
        existing.put(message)
        writeJson(historyName(threadId), JSONObject().put("messages", existing))
    }

    /** Every thread this device holds, group or person. */
    fun threads(): List<String> =
        storage.names().filter { it.startsWith(HISTORY_PREFIX) }
            .map { it.removePrefix(HISTORY_PREFIX) }

    private fun historyName(threadId: String) = HISTORY_PREFIX + threadId

    // -- Forgetting ----------------------------------------------------------

    /** Forget one device's session. The identity and every other session survive. */
    fun forget(peerDeviceId: String) = storage.delete(sessionName(peerDeviceId))

    /**
     * Delete a thread's history.
     *
     * ⚠️ Deliberately does NOT delete the session. Clearing what was said is
     * not the same as ending the conversation, and dropping the ratchet would
     * make the next message from that person undecryptable — a "clear chat"
     * that quietly breaks the chat.
     */
    fun forgetHistory(threadId: String) = storage.delete(historyName(threadId))

    /**
     * Forget everything, identity included.
     *
     * This is sign-out and account deletion. Leaving the identity behind would
     * mean the next person to use the phone inherits the device the relay still
     * has published — and messages addressed to it.
     */
    fun forgetEverything() = storage.names().forEach(storage::delete)

    /** Forget conversations but keep this device's identity. */
    fun forgetAllSessions() =
        storage.names().filter { it.startsWith(SESSION_PREFIX) }.forEach(storage::delete)
}