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? {
        val salt = unhex(browser.salt)
        val key = VaultCrypto.deriveKey(password, salt)
        try {
            val sealed = unhex(browser.nonce) + unhex(browser.ciphertext) + unhex(browser.tag)
            val 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()
            val payload = VaultCodec.decode(plain.toByteArray(Charsets.UTF_8))

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

            return StoredVault(
                version = 1,
                salt = b64(salt),
                blob = b64(VaultCrypto.encrypt(VaultCodec.encode(migrated), key, VaultCrypto.AAD_SLOT_A)),
                rootKeyEscrow = browser.toRootKeyEscrow(),
            )
        } finally {
            VaultCrypto.zero(key)
        }
    }

    /** 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
        return runCatching {
            val hexBytes = VaultCrypto.decrypt(
                unhex(nonce) + unhex(ct) + unhex(tag),
                derivedKey,
                VaultCrypto.AAD_ROOT_KEY,
            )
            unhex(hexBytes.toString(Charsets.UTF_8).trimBucketPadding())
        }.getOrNull()
    }

    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 = b64(unhex(s)),
            wrappedRootKey = b64(unhex(w)),
            nonce = b64(unhex(n)),
            tag = b64(unhex(t)),
        )
    }

    private 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")
    }

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

/**
 * 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() }