package com.peachpasswords.vault

import kotlinx.serialization.SerialName
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json
import java.util.Base64

/**
 * Reads the browser extension's slot A so Android can restore a vault that was
 * created there — a Codex printed from the extension, most importantly.
 *
 * The two platforms already agree on everything that matters cryptographically:
 * Argon2id at m=65536/t=3/p=4 over the NFC-normalized password, AES-256-GCM
 * with AAD `slot_a_v1`, and a slot-A plaintext that is plain JSON
 * (`{entries, duressConfigured}`) which Android's [VaultCodec] parses as-is.
 * What differs is only how the sealed bytes are written down:
 *
 *   browser   { salt, nonce, ciphertext, tag }   all lowercase hex, separate
 *   Android   { salt, blob }                     base64, blob = nonce||ct||tag
 *
 * Android's [VaultStore.decodeSlot] required the Android shape, so a real
 * browser Codex decoded all 1978 of its Aztec codes and then died at the last
 * step with "invalid_codex_blob" — the vault was intact and readable the whole
 * time, sitting behind a field-name mismatch.
 *
 * The root key needs moving too, and is the reason this is a translation rather
 * than a rename: the browser keeps it wrapped in the slot (`pw_root_*`, AAD
 * `pw_root_v1`, plaintext = lowercase hex of the 32 bytes) while Android keeps
 * it inside the encrypted payload as `rootKeyB64`. Carrying it across is what
 * makes the restored vault keep the same recovery words and the same sync
 * identity it had on the browser; minting a fresh one would silently invalidate
 * the words the user has written down.
 */
@Serializable
data class BrowserSlotA(
    val version: Int = 1,
    val salt: String,
    val nonce: String,
    val ciphertext: String,
    val tag: String,
    @SerialName("recovery_salt") val recoverySalt: String? = null,
    @SerialName("recovery_wrapped_root_key") val recoveryWrappedRootKey: String? = null,
    @SerialName("recovery_nonce") val recoveryNonce: String? = null,
    @SerialName("recovery_tag") val recoveryTag: String? = null,
    @SerialName("pw_root_ciphertext") val pwRootCiphertext: String? = null,
    @SerialName("pw_root_nonce") val pwRootNonce: String? = null,
    @SerialName("pw_root_tag") val pwRootTag: String? = null,
)

object BrowserSlot {
    private val json = Json { ignoreUnknownKeys = true; isLenient = true }

    /** Null when [s] is not a browser slot — callers try the Android shape first. */
    fun parse(s: String): BrowserSlotA? =
        runCatching { json.decodeFromString<BrowserSlotA>(s) }.getOrNull()

    /**
     * Re-seal a browser slot as a native Android one under the same password.
     *
     * Returns null on a wrong password (the GCM tag fails), which the caller
     * must surface as a wrong-password error rather than a corrupt-blob one —
     * they are very different things to a user holding a paper Codex.
     */
    fun toAndroidSlot(browser: BrowserSlotA, password: String): StoredVault? {
        var salt: ByteArray? = null
        var key: ByteArray? = null
        var sealed: ByteArray? = null
        var payloadJson: ByteArray? = null
        var plainBytes: ByteArray? = null
        var encoded: ByteArray? = null
        var encrypted: ByteArray? = null
        var rootKey: ByteArray? = null
        try {
            salt = unhex(browser.salt)
            key = VaultCrypto.deriveKey(password, salt!!)
            sealed = sealedHex(browser.nonce, browser.ciphertext, browser.tag)
            payloadJson = try {
                VaultCrypto.decrypt(sealed!!, key!!, VaultCrypto.AAD_SLOT_A)
            } catch (_: Exception) {
                return null
            }
            // The browser pads its plaintext up to a size bucket with ZERO
            // BYTES (`padToNearestBucket` in peachPayload.ts — "zero-filled in
            // the padding region"), to stop file size leaking vault size.
            //
            // This said "trailing SPACES" and used a plain trim(), which only
            // strips whitespace — and NUL is not whitespace to Kotlin. So the
            // padding survived and JSON parsing died the moment the real
            // payload ended, reporting an offset equal to the payload's own
            // length. It looked like a corrupt Codex; the vault was fine and
            // the password was right.
            val plain = payloadJson!!.toString(Charsets.UTF_8).trimBucketPadding()
            plainBytes = plain.toByteArray(Charsets.UTF_8)
            val payload = VaultCodec.decode(plainBytes!!)

            rootKey = unwrapRootKey(browser, key!!)
            val migrated = payload.copy(
                rootKeyB64 = rootKey?.let { b64(it) } ?: payload.rootKeyB64,
            )

            encoded = VaultCodec.encode(migrated)
            encrypted = VaultCrypto.encrypt(encoded!!, key!!, VaultCrypto.AAD_SLOT_A)

            return StoredVault(
                version = 1,
                salt = b64(salt!!),
                blob = b64(encrypted!!),
                rootKeyEscrow = browser.toRootKeyEscrow(),
            )
        } finally {
            encrypted?.let(VaultCrypto::zero)
            encoded?.let(VaultCrypto::zero)
            plainBytes?.let(VaultCrypto::zero)
            payloadJson?.let(VaultCrypto::zero)
            sealed?.let(VaultCrypto::zero)
            rootKey?.let(VaultCrypto::zero)
            salt?.let(VaultCrypto::zero)
            key?.let(VaultCrypto::zero)
        }
    }

    /** The browser wraps the root key's lowercase hex, not its raw bytes. */
    private fun unwrapRootKey(browser: BrowserSlotA, derivedKey: ByteArray): ByteArray? {
        val ct = browser.pwRootCiphertext ?: return null
        val nonce = browser.pwRootNonce ?: return null
        val tag = browser.pwRootTag ?: return null
        var sealed: ByteArray? = null
        var hexBytes: ByteArray? = null
        return try {
            sealed = sealedHex(nonce, ct, tag)
            hexBytes = VaultCrypto.decrypt(
                sealed!!,
                derivedKey,
                VaultCrypto.AAD_ROOT_KEY,
            )
            unhex(hexBytes!!.toString(Charsets.UTF_8).trimBucketPadding())
        } catch (_: Exception) {
            null
        } finally {
            hexBytes?.let(VaultCrypto::zero)
            sealed?.let(VaultCrypto::zero)
        }
    }

    private fun BrowserSlotA.toRootKeyEscrow(): RootKeyEscrow? {
        val s = recoverySalt ?: return null
        val w = recoveryWrappedRootKey ?: return null
        val n = recoveryNonce ?: return null
        val t = recoveryTag ?: return null
        return RootKeyEscrow(
            salt = b64Hex(s),
            wrappedRootKey = b64Hex(w),
            nonce = b64Hex(n),
            tag = b64Hex(t),
        )
    }

    /** Decode concatenated browser hex fields while scrubbing each component. */
    private fun sealedHex(vararg parts: String): ByteArray {
        val decoded = ArrayList<ByteArray>(parts.size)
        return try {
            parts.forEach { decoded += unhex(it) }
            val out = ByteArray(decoded.sumOf { it.size })
            var offset = 0
            decoded.forEach {
                System.arraycopy(it, 0, out, offset, it.size)
                offset += it.size
            }
            out
        } finally {
            decoded.forEach(VaultCrypto::zero)
        }
    }

    private fun b64Hex(hex: String): String {
        val bytes = unhex(hex)
        return try {
            b64(bytes)
        } finally {
            VaultCrypto.zero(bytes)
        }
    }

    private fun b64(b: ByteArray): String = Base64.getEncoder().encodeToString(b)
}

/** Hex -> bytes. Lowercase or uppercase digits; throws on odd length or bad char. */
internal fun unhex(s: String): ByteArray {
    require(s.length % 2 == 0) { "odd-length hex" }
    return ByteArray(s.length / 2) { i ->
        ((hexDigit(s[i * 2]) shl 4) or hexDigit(s[i * 2 + 1])).toByte()
    }
}

private fun hexDigit(c: Char): Int = when (c) {
    in '0'..'9' -> c - '0'
    in 'a'..'f' -> c - 'a' + 10
    in 'A'..'F' -> c - 'A' + 10
    else -> throw IllegalArgumentException("bad hex digit")
}

/**
 * Strip the browser's size-bucket padding: zero bytes, plus any surrounding
 * whitespace. `String.trim()` alone is not enough — `Char.isWhitespace()` is
 * false for NUL, so the padding sails straight into the JSON parser.
 */
private fun String.trimBucketPadding(): String = trim { it == '\u0000' || it.isWhitespace() }