package io.flamenet.e2e

import org.json.JSONArray
import org.json.JSONObject

/**
 * This device's history, translated into the shape both apps agree on.
 *
 * Section 16 pins the envelope: the KDF, the canonical header used as AAD and
 * the AES-GCM sealing. It says nothing about what is sealed. Until now each app
 * sealed its own storage verbatim and stamped the file with the platform that
 * wrote it, and each then refused the other's, in as many words: "import it on
 * a device of the same kind". Somebody moving between an Android phone and an
 * iPhone lost every conversation they had.
 *
 * The iPhone half of the shared shape landed first. This is the other half.
 *
 * It is a translation and not a dump, which is the whole point. This app stores
 * a message under the keys its wire format uses (`secs`, `fwd`, `rx`), keeps
 * one-to-one and group threads in a single map under prefixed ids, and holds
 * expiries in unix seconds. The shared shape spells those out, separates the
 * two kinds of thread and counts in milliseconds. Exporting the storage as-is
 * is exactly what made the two formats incompatible in the first place.
 *
 * Three rules hold it together.
 *
 * **Times are epoch milliseconds, as integers.** This app already keeps `at` in
 * millis, so messages pass straight through; expiries do not, and are converted
 * at the boundary. A backup that disagreed about the unit would restore every
 * story at the wrong time without ever failing.
 *
 * **A key this file does not name is not carried.** Both sides build a fresh
 * object rather than copying, so adding a field to storage and expecting it to
 * appear in a backup does not work. `BackupPayloadTest` reads this source and
 * fails if a known key stops being mentioned.
 *
 * **What the other app cannot represent is kept apart rather than faked.**
 * Federated threads are the case that matters: this app keys them by address,
 * and the iPhone's history has no room for one, so they travel under their own
 * key. An iPhone ignores it and an Android phone gets them back.
 */
object BackupPayload {
    const val VERSION = 1

    /** What a section 16 backup is called on disk. Shared with the iPhone. */
    const val FILE_EXTENSION = "fnbackup"

    private const val GROUP_PREFIX = "group:"
    private const val ADDRESS_PREFIX = "addr:"

    /** Federated threads, which the iPhone has nowhere to put. */
    private const val ADDRESS_THREADS = "addrThreads"

    // -- Writing -------------------------------------------------------------.

    /** Turn [internal], as [Vault.exportPayload] produces it, into the shared shape. */
    fun toShared(internal: JSONObject): JSONObject {
        val root = JSONObject().put("v", VERSION)

        val threads = JSONObject()
        val groups = JSONObject()
        val addressed = JSONObject()
        internal.optJSONObject("threads")?.let { all ->
            for (tid in all.keys()) {
                val messages = all.optJSONArray(tid) ?: continue
                when {
                    tid.startsWith(GROUP_PREFIX) -> {
                        groups.put(tid.removePrefix(GROUP_PREFIX), mapArray(messages, ::messageToShared))
                    }

                    tid.startsWith(ADDRESS_PREFIX) -> {
                        // The prefix comes OFF, exactly as the group one does
                        // above. It is how this app namespaces one `threads`
                        // map; the shared shape keys these by the address
                        // itself. Leaving it on wrote `addr:ada@relay.example`
                        // into the file, which is not an address, and the
                        // iPhone writes the bare one, so neither side could
                        // read the other's federated conversations.
                        addressed.put(
                            tid.removePrefix(ADDRESS_PREFIX),
                            mapArray(messages, ::messageToShared),
                        )
                    }

                    else -> {
                        threads.put(tid, mapArray(messages, ::messageToShared))
                    }
                }
            }
        }
        if (threads.length() > 0) root.put("threads", threads)
        if (groups.length() > 0) root.put("groups", groups)
        if (addressed.length() > 0) root.put(ADDRESS_THREADS, addressed)

        internal.optJSONArray("callLog")?.takeIf { it.length() > 0 }?.let { log ->
            root.put(
                "calls",
                mapArray(log) { c ->
                    JSONObject()
                        .put("peer", c.optLong("peer"))
                        .put("at", c.optLong("at"))
                        .put("in", c.optBoolean("in"))
                        // The only rename here: this app writes `ok`.
                        .put("connected", c.optBoolean("ok"))
                },
            )
        }

        // Stories are one flat list here and are keyed by author over there.
        internal.optJSONObject("stories")?.optJSONArray("items")?.let { items ->
            val byAuthor = JSONObject()
            for (i in 0 until items.length()) {
                val s = items.optJSONObject(i) ?: continue
                val author = s.optLong("author").toString()
                val list = byAuthor.optJSONArray(author) ?: JSONArray().also { byAuthor.put(author, it) }
                list.put(storyToShared(s))
            }
            if (byAuthor.length() > 0) root.put("stories", byAuthor)
        }

        internal.optJSONObject("roster")?.let { r ->
            r.optJSONArray("items")?.takeIf { it.length() > 0 }?.let { root.put("roster", it) }
            if (r.has("offServer")) root.put("rosterOffServer", r.optBoolean("offServer"))
        }

        internal.optJSONObject("groupNames")?.takeIf { it.length() > 0 }?.let { root.put("groupNames", it) }

        internal.optJSONObject("groupSecrets")?.takeIf { it.length() > 0 }?.let { all ->
            val out = JSONObject()
            for (gid in all.keys()) {
                val s = all.optJSONObject(gid) ?: continue
                val one =
                    JSONObject()
                        .put("tag", s.optString("tag"))
                        .put("key", s.optString("key"))
                        // The relay refuses a roster rewrite whose version does
                        // not increase, so a restore that reset it would lock
                        // the group out of its own next rekey.
                        .put("version", s.optInt("version"))
                s.optString("owner").takeIf { it.isNotEmpty() }?.let { one.put("ownerPrivate", it) }
                out.put(gid, one)
            }
            root.put("groupSecrets", out)
        }

        // Who reacted to which story, keyed by story then by reactor. The
        // iPhone writes and reads the same shape, so a tally survives a move
        // between the two.
        internal.optJSONObject("storyReactions")?.takeIf { it.length() > 0 }?.let {
            root.put("storyReactions", it)
        }

        // §16.3. State somebody set on purpose, under the iPhone's names.
        internal.optJSONObject("storyAudience")?.let { a ->
            // This app stores the enum's NAME, which is upper case; the iPhone
            // writes lower case and reads nothing else, so an audience crossing
            // between the two would have been dropped in silence and the person
            // would have restored a phone that shows their stories to everyone.
            root.put(
                "storyAudience",
                JSONObject()
                    .put("mode", a.optString("mode").lowercase())
                    .put("ids", a.optJSONArray("ids") ?: JSONArray()),
            )
        }
        internal.optJSONObject("storyMutes")?.optJSONArray("ids")?.takeIf { it.length() > 0 }?.let {
            root.put("mutedStoryAuthors", it)
        }
        internal.optJSONObject("storyArchive")?.optJSONArray("items")?.let { items ->
            val out = JSONArray()
            for (i in 0 until items.length()) {
                items.optJSONObject(i)?.let { out.put(storyToShared(it)) }
            }
            if (out.length() > 0) root.put("storyArchive", out)
        }
        internal.optJSONObject("drafts")?.takeIf { it.length() > 0 }?.let { root.put("drafts", it) }
        // Carried whole, `at` and `by` included, though this app shows neither.
        // Dropping a field an engine cannot use turns a backup into a lossy
        // hop: restore an iPhone through an Android phone and the pins come
        // back stripped of who pinned them and when.
        internal.optJSONObject("pins")?.takeIf { it.length() > 0 }?.let { root.put("pins", it) }
        internal.optJSONObject("polls")?.takeIf { it.length() > 0 }?.let { root.put("polls", it) }
        internal.optJSONObject("presenceTokensIssued")?.takeIf { it.length() > 0 }?.let {
            root.put("presenceTokensIssued", it)
        }
        internal.optJSONObject("presenceTokensHeld")?.takeIf { it.length() > 0 }?.let {
            root.put("presenceTokensHeld", it)
        }

        return root
    }

    private fun messageToShared(m: JSONObject): JSONObject {
        val out =
            JSONObject()
                .put("mine", m.optBoolean("mine"))
                .put("at", m.optLong("at"))
                .put("text", m.optString("text"))
        m.optString("mid").takeIf { it.isNotEmpty() }?.let { out.put("mid", it) }
        // The author rides along on group threads, where it says who spoke.
        // Only when there is one: this app writes 0 on a one-to-one message,
        // where the author is whoever is not you, and emitting that would put
        // a claim in the file that the message came from user zero.
        m.optLong("author", 0L).takeIf { it != 0L }?.let { out.put("author", it) }
        m.optJSONObject("att")?.let {
            out.put("att", attachmentToShared(it, m.optString("kind", Attachment.IMAGE)))
        }
        m.optJSONObject("rx")?.takeIf { it.length() > 0 }?.let { out.put("reactions", it) }
        m.optLong("edited", 0L).takeIf { it > 0L }?.let { out.put("edited", it) }
        // Both shapes, for the reason `Content.isForwarded` accepts both:
        // `optBoolean` does not coerce a number, so a `"fwd":1` from another
        // engine would read as false.
        if (m.optBoolean("fwd", false) || m.optInt("fwd", 0) != 0) out.put("forwarded", true)
        m.optInt("secs", 0).takeIf { it > 0 }?.let { out.put("lifetime", it) }
        m.optJSONObject("re")?.let { out.put("re", it) }
        m.optJSONObject("stk")?.let { out.put("stk", it) }
        m.optJSONObject("lnk")?.let { out.put("lnk", it) }
        m.optString("delivery").takeIf { it.isNotEmpty() }?.let { out.put("delivery", it.lowercase()) }
        return out
    }

    /**
     * The kind lives beside `att` here and inside it over there, and the sizes
     * and duration are spelled out rather than abbreviated.
     */
    private fun attachmentToShared(
        a: JSONObject,
        kind: String,
    ): JSONObject {
        val out =
            JSONObject()
                .put("kind", kind)
                .put("id", a.optString("id"))
                .put("key", a.optString("key"))
                .put("digest", a.optString("digest"))
                .put("bytes", a.optInt("bytes"))
                .put("mime", a.optString("mime"))
                .put("width", a.optInt("w"))
                .put("height", a.optInt("h"))
        a.optDouble("dur", 0.0).takeIf { it > 0.0 }?.let { out.put("duration", it) }
        // Seconds here, milliseconds there.
        a.optDouble("exp", 0.0).takeIf { it > 0.0 }?.let { out.put("expires", secondsToMillis(it)) }
        // Stored under the same names the shared shape uses, so these three
        // pass straight through. They were not carried before because this app
        // did not keep them at all; a restored document arrived with no name
        // and no label, which is worse than not restoring it.
        a.optString("filename").takeIf { it.isNotEmpty() }?.let { out.put("filename", it) }
        a.optInt("classification", 0).takeIf { it != 0 }?.let { out.put("classification", it) }
        a.optJSONArray("compartments")?.takeIf { it.length() > 0 }?.let { out.put("compartments", it) }
        return out
    }

    /**
     * A story in the shared shape.
     *
     * `key` is the identity and `id` is NOT written. The two engines had the
     * same idea and put it in different places: this app keeps the shared
     * identity straight in `id`, while the iPhone numbers a story locally and
     * keeps the identity in `key`. One name, two meanings, so a story
     * travelling between the two came back under a different id and imported
     * as a second copy of itself.
     */
    private fun storyToShared(s: JSONObject): JSONObject {
        val out =
            JSONObject()
                .put("author", s.optLong("author"))
                .put("mine", s.optBoolean("mine"))
                .put("caption", s.optString("caption"))
                .put("at", s.optLong("at"))
        s.optDouble("expires", 0.0).takeIf { it > 0.0 }?.let { out.put("expires", secondsToMillis(it)) }
        s.optJSONObject("ref")?.let { out.put("att", attachmentToShared(it, Attachment.STORY)) }
        s.optString("gid").takeIf { it.isNotEmpty() }?.let { out.put("gid", it) }
        s.optString("stream").takeIf { it.isNotEmpty() }?.let { out.put("stream", it) }
        // Always written. This app's `id` IS the identity, so it is the
        // fallback when an older record carried no separate key, and the post
        // time is the last resort rather than letting a story travel with no
        // way to recognise it.
        out.put(
            "key",
            s.optString("key").ifEmpty { s.optString("id") }.ifEmpty { s.optLong("at").toString() },
        )
        s.optJSONObject("set")?.let { out.put("set", it) }
        // Whether the strip has been watched. This used to be ours alone and
        // travelled anyway, on the grounds that it was the difference between a
        // strip remembering what you watched and one showing everything as new.
        // The iPhone turned out to HAVE the mark and not to be writing it, so
        // it relit every ring at each launch; both engines carry it now and it
        // is in the canonical set rather than a divergence.
        if (s.optBoolean("seen", false)) out.put("seen", true)
        return out
    }

    // -- Reading -------------------------------------------------------------.

    class NewerBackup(
        val version: Int,
    ) : Exception("This backup was written by a newer version of the app.")

    /** Turn a shared payload back into what [Vault.mergeImportedPayload] expects. */
    fun fromShared(shared: JSONObject): JSONObject {
        val version = shared.optInt("v", 0)
        if (version > VERSION) throw NewerBackup(version)

        val threads = JSONObject()
        shared.optJSONObject("threads")?.let { all ->
            for (tid in all.keys()) {
                all.optJSONArray(tid)?.let { threads.put(tid, mapArray(it, ::messageFromShared)) }
            }
        }
        shared.optJSONObject("groups")?.let { all ->
            for (gid in all.keys()) {
                all.optJSONArray(gid)?.let {
                    threads.put(GROUP_PREFIX + gid, mapArray(it, ::messageFromShared))
                }
            }
        }
        shared.optJSONObject(ADDRESS_THREADS)?.let { all ->
            for (address in all.keys()) {
                // The prefix goes back on. A file written before it was
                // stripped carries it already, and adding a second one would
                // make a thread nobody can find.
                val tid =
                    if (address.startsWith(ADDRESS_PREFIX)) address else ADDRESS_PREFIX + address
                all.optJSONArray(address)?.let {
                    threads.put(tid, mapArray(it, ::messageFromShared))
                }
            }
        }

        val out = JSONObject().put("threads", threads)

        shared.optJSONArray("calls")?.takeIf { it.length() > 0 }?.let { calls ->
            out.put(
                "callLog",
                mapArray(calls) { c ->
                    JSONObject()
                        .put("peer", c.optLong("peer"))
                        .put("at", c.optLong("at"))
                        .put("in", c.optBoolean("in"))
                        .put("ok", c.optBoolean("connected"))
                },
            )
        }

        shared.optJSONObject("stories")?.let { byAuthor ->
            val items = JSONArray()
            for (author in byAuthor.keys()) {
                val list = byAuthor.optJSONArray(author) ?: continue
                for (i in 0 until list.length()) {
                    val s = list.optJSONObject(i) ?: continue
                    items.put(storyFromShared(s, author.toLongOrNull() ?: 0L))
                }
            }
            if (items.length() > 0) out.put("stories", JSONObject().put("items", items))
        }

        shared.optJSONArray("roster")?.takeIf { it.length() > 0 }?.let {
            out.put(
                "roster",
                JSONObject().put("items", it).put("offServer", shared.optBoolean("rosterOffServer", false)),
            )
        }

        shared.optJSONObject("groupNames")?.takeIf { it.length() > 0 }?.let { out.put("groupNames", it) }

        shared.optJSONObject("groupSecrets")?.takeIf { it.length() > 0 }?.let { all ->
            val secrets = JSONObject()
            for (gid in all.keys()) {
                val s = all.optJSONObject(gid) ?: continue
                val one =
                    JSONObject()
                        .put("tag", s.optString("tag"))
                        .put("key", s.optString("key"))
                        .put("version", s.optInt("version"))
                s.optString("ownerPrivate").takeIf { it.isNotEmpty() }?.let { one.put("owner", it) }
                secrets.put(gid, one)
            }
            out.put("groupSecrets", secrets)
        }

        shared.optJSONObject("storyReactions")?.takeIf { it.length() > 0 }?.let {
            out.put("storyReactions", it)
        }

        shared.optJSONObject("storyAudience")?.let { a ->
            out.put(
                "storyAudience",
                JSONObject()
                    .put("mode", a.optString("mode").uppercase())
                    .put("ids", a.optJSONArray("ids") ?: JSONArray()),
            )
        }
        shared.optJSONArray("mutedStoryAuthors")?.takeIf { it.length() > 0 }?.let {
            out.put("storyMutes", JSONObject().put("ids", it))
        }
        shared.optJSONArray("storyArchive")?.takeIf { it.length() > 0 }?.let { items ->
            val restored = JSONArray()
            for (i in 0 until items.length()) {
                val s = items.optJSONObject(i) ?: continue
                restored.put(storyFromShared(s, s.optLong("author", 0L)))
            }
            if (restored.length() > 0) out.put("storyArchive", JSONObject().put("items", restored))
        }
        shared.optJSONObject("drafts")?.takeIf { it.length() > 0 }?.let { out.put("drafts", it) }
        shared.optJSONObject("pins")?.takeIf { it.length() > 0 }?.let { out.put("pins", it) }
        shared.optJSONObject("polls")?.takeIf { it.length() > 0 }?.let { out.put("polls", it) }
        shared.optJSONObject("presenceTokensIssued")?.takeIf { it.length() > 0 }?.let {
            out.put("presenceTokensIssued", it)
        }
        shared.optJSONObject("presenceTokensHeld")?.takeIf { it.length() > 0 }?.let {
            out.put("presenceTokensHeld", it)
        }

        return out
    }

    private fun messageFromShared(m: JSONObject): JSONObject {
        val out =
            JSONObject()
                .put("mine", m.optBoolean("mine"))
                .put("at", m.optLong("at"))
                .put("text", m.optString("text"))
        // Storage keeps an author on every message and reads a missing one as
        // 0, so this is written only when the file carried one.
        m.optLong("author", 0L).takeIf { it != 0L }?.let { out.put("author", it) }
        m.optString("mid").takeIf { it.isNotEmpty() }?.let { out.put("mid", it) }
        m.optJSONObject("att")?.let { att ->
            out.put("kind", att.optString("kind", Attachment.IMAGE))
            out.put("att", attachmentFromShared(att))
        }
        m.optJSONObject("reactions")?.takeIf { it.length() > 0 }?.let { out.put("rx", it) }
        m.optLong("edited", 0L).takeIf { it > 0L }?.let { out.put("edited", it) }
        if (m.optBoolean("forwarded", false) || m.optInt("forwarded", 0) != 0) out.put("fwd", true)
        m.optInt("lifetime", 0).takeIf { it > 0 }?.let { out.put("secs", it) }
        m.optJSONObject("re")?.let { out.put("re", it) }
        m.optJSONObject("stk")?.let { out.put("stk", it) }
        m.optJSONObject("lnk")?.let { out.put("lnk", it) }
        // Stored upper case here, lower case on the wire and in the file.
        m.optString("delivery").takeIf { it.isNotEmpty() }?.let { out.put("delivery", it.uppercase()) }
        return out
    }

    private fun attachmentFromShared(a: JSONObject): JSONObject {
        val out =
            JSONObject()
                .put("id", a.optString("id"))
                .put("key", a.optString("key"))
                .put("digest", a.optString("digest"))
                .put("bytes", a.optInt("bytes"))
                .put("mime", a.optString("mime"))
                .put("w", a.optInt("width"))
                .put("h", a.optInt("height"))
        a.optDouble("duration", 0.0).takeIf { it > 0.0 }?.let { out.put("dur", it) }
        a.optDouble("expires", 0.0).takeIf { it > 0.0 }?.let { out.put("exp", millisToSeconds(it)) }
        a.optString("filename").takeIf { it.isNotEmpty() }?.let { out.put("filename", it) }
        a.optInt("classification", 0).takeIf { it != 0 }?.let { out.put("classification", it) }
        a.optJSONArray("compartments")?.takeIf { it.length() > 0 }?.let { out.put("compartments", it) }
        return out
    }

    private fun storyFromShared(
        s: JSONObject,
        author: Long,
    ): JSONObject {
        val out =
            JSONObject()
                // The identity, which this app stores as its `id`. `key` is
                // where it lives in the shared shape; a file written before
                // that was settled carries it in `id`, and the post time is the
                // last resort. `mergeImportedPayload` de-duplicates on this
                // string, so it only has to be stable, not numeric.
                .put(
                    "id",
                    s
                        .optString("key")
                        // `id` only when it is a STRING. A number there is the
                        // writing device's own row handle and names nothing
                        // here; adopting it is how a story imported as a second
                        // copy of itself. The post time is the last resort, and
                        // it is what the iPhone falls back to as well, so the
                        // two still land on the same value.
                        .ifEmpty { (s.opt("id") as? String).orEmpty() }
                        .ifEmpty { s.optLong("at").toString() },
                )
                .put("author", if (s.has("author")) s.optLong("author") else author)
                .put("mine", s.optBoolean("mine"))
                .put("caption", s.optString("caption"))
                .put("at", s.optLong("at"))
                .put("seen", s.optBoolean("seen", false))
        s.optDouble("expires", 0.0).takeIf { it > 0.0 }?.let { out.put("expires", millisToSeconds(it)) }
        // `ref` carries the expiry as well, because that is where the strip
        // reads it back from when deciding whether the bytes may still render.
        s.optJSONObject("att")?.let { out.put("ref", attachmentFromShared(it)) }
        s.optString("gid").takeIf { it.isNotEmpty() }?.let { out.put("gid", it) }
        s.optString("stream").takeIf { it.isNotEmpty() }?.let { out.put("stream", it) }
        s.optJSONObject("set")?.let { out.put("set", it) }
        return out
    }

    // -- Small shared pieces -------------------------------------------------.

    private fun secondsToMillis(seconds: Double): Long = Math.round(seconds * 1000.0)

    private fun millisToSeconds(millis: Double): Double = millis / 1000.0

    private fun mapArray(
        source: JSONArray,
        transform: (JSONObject) -> JSONObject,
    ): JSONArray {
        val out = JSONArray()
        for (i in 0 until source.length()) {
            source.optJSONObject(i)?.let { out.put(transform(it)) }
        }
        return out
    }
}