package no.zyon.shade.serialization

import no.zyon.shade.types.ChainState
import no.zyon.shade.types.IdentityKeyPair
import no.zyon.shade.types.KeyPair
import no.zyon.shade.types.OneTimePreKey
import no.zyon.shade.types.SessionState
import no.zyon.shade.types.SignedPreKey
import org.json.JSONObject
import java.util.Base64 as JdkBase64

/**
 * Plain-JSON serialization for the persisted protocol state types
 * (`IdentityKeyPair`, `SignedPreKey`, `OneTimePreKey`, `SessionState`).
 *
 * The on-disk shape is for at-rest storage only — it does NOT need
 * to round-trip across platforms (TS uses its own JSON shape via
 * `@shade/core/serialization`). What matters is that the Kotlin
 * round-trip (`serialize` then `deserialize`) preserves every byte.
 *
 * Both Android-targeted (`shade-android-keystore`) and pure-JVM
 * (`shade-android` tests) callers use this — the function works
 * without any `android.*` imports so it compiles in both.
 */
object SessionStateJson {

    fun serialize(state: SessionState): String {
        val o = JSONObject()
        o.put("remoteIdentityKey", b64(state.remoteIdentityKey))
        o.put("rootKey", b64(state.rootKey))
        o.put("sendChain", chainToJson(state.sendChain))
        if (state.receiveChain != null) o.put("receiveChain", chainToJson(state.receiveChain!!))
        o.put("dhSend", keyPairToJson(state.dhSend))
        if (state.dhReceive != null) o.put("dhReceive", b64(state.dhReceive!!))
        o.put("previousSendCounter", state.previousSendCounter)
        val skipped = JSONObject()
        for ((k, v) in state.skippedKeys) skipped.put(k, b64(v))
        o.put("skippedKeys", skipped)
        return o.toString()
    }

    fun deserialize(s: String): SessionState {
        val o = JSONObject(s)
        val skipped = mutableMapOf<String, ByteArray>()
        val skJson = o.optJSONObject("skippedKeys")
        if (skJson != null) {
            val it = skJson.keys()
            while (it.hasNext()) {
                val k = it.next()
                skipped[k] = fb64(skJson.getString(k))
            }
        }
        return SessionState(
            remoteIdentityKey = fb64(o.getString("remoteIdentityKey")),
            rootKey = fb64(o.getString("rootKey")),
            sendChain = chainFromJson(o.getJSONObject("sendChain")),
            receiveChain = if (o.has("receiveChain"))
                chainFromJson(o.getJSONObject("receiveChain")) else null,
            dhSend = keyPairFromJson(o.getJSONObject("dhSend")),
            dhReceive = if (o.has("dhReceive")) fb64(o.getString("dhReceive")) else null,
            previousSendCounter = o.getInt("previousSendCounter"),
            skippedKeys = skipped,
        )
    }

    fun serializeIdentityKeyPair(k: IdentityKeyPair): String = JSONObject().apply {
        put("signingPublicKey", b64(k.signingPublicKey))
        put("signingPrivateKey", b64(k.signingPrivateKey))
        put("dhPublicKey", b64(k.dhPublicKey))
        put("dhPrivateKey", b64(k.dhPrivateKey))
    }.toString()

    fun deserializeIdentityKeyPair(s: String): IdentityKeyPair = JSONObject(s).run {
        IdentityKeyPair(
            signingPublicKey = fb64(getString("signingPublicKey")),
            signingPrivateKey = fb64(getString("signingPrivateKey")),
            dhPublicKey = fb64(getString("dhPublicKey")),
            dhPrivateKey = fb64(getString("dhPrivateKey")),
        )
    }

    fun serializeSignedPreKey(k: SignedPreKey): String = JSONObject().apply {
        put("keyId", k.keyId)
        put("keyPair", keyPairToJson(k.keyPair))
        put("signature", b64(k.signature))
        put("timestamp", k.timestamp)
    }.toString()

    fun deserializeSignedPreKey(s: String): SignedPreKey = JSONObject(s).run {
        SignedPreKey(
            keyId = getInt("keyId"),
            keyPair = keyPairFromJson(getJSONObject("keyPair")),
            signature = fb64(getString("signature")),
            timestamp = getLong("timestamp"),
        )
    }

    fun serializeOneTimePreKey(k: OneTimePreKey): String = JSONObject().apply {
        put("keyId", k.keyId)
        put("keyPair", keyPairToJson(k.keyPair))
    }.toString()

    fun deserializeOneTimePreKey(s: String): OneTimePreKey = JSONObject(s).run {
        OneTimePreKey(
            keyId = getInt("keyId"),
            keyPair = keyPairFromJson(getJSONObject("keyPair")),
        )
    }

    private fun chainToJson(c: ChainState): JSONObject = JSONObject().apply {
        put("chainKey", b64(c.chainKey))
        put("counter", c.counter)
    }

    private fun chainFromJson(o: JSONObject): ChainState =
        ChainState(chainKey = fb64(o.getString("chainKey")), counter = o.getInt("counter"))

    private fun keyPairToJson(k: KeyPair): JSONObject = JSONObject().apply {
        put("publicKey", b64(k.publicKey))
        put("privateKey", b64(k.privateKey))
    }

    private fun keyPairFromJson(o: JSONObject): KeyPair = KeyPair(
        publicKey = fb64(o.getString("publicKey")),
        privateKey = fb64(o.getString("privateKey")),
    )

    // android.util.Base64 isn't on the JVM classpath; java.util.Base64
    // is available on both modern JVM and Android API 26+. Use JDK
    // Base64 throughout — it's present on both targets.
    private fun b64(b: ByteArray): String = JdkBase64.getEncoder().encodeToString(b)
    private fun fb64(s: String): ByteArray = JdkBase64.getDecoder().decode(s)
}
