package com.peachpasswords.vault

import org.bouncycastle.crypto.generators.Argon2BytesGenerator
import org.bouncycastle.crypto.params.Argon2Parameters
import java.security.SecureRandom
import java.text.Normalizer
import javax.crypto.Cipher
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.SecretKeySpec

object VaultCrypto {
    const val ARGON2_MEMORY_KB = 65536
    const val ARGON2_ITERATIONS = 3
    const val ARGON2_PARALLELISM = 4
    const val ARGON2_HASH_LENGTH = 32
    const val SALT_SIZE = 16
    const val NONCE_SIZE = 12
    const val TAG_LENGTH_BITS = 128
    const val KEY_SIZE = 32
    const val PAD_BLOCK = 1024
    const val DURESS_SLOT_SIZE = 262_144

    val PADDING_BUCKETS = intArrayOf(1024, 4096, 16384, 65536, 262144)

    const val AAD_SLOT_A = "slot_a_v1"
    const val AAD_SLOT_B = "slot_b_v1"
    const val AAD_ROOT_KEY = "pw_root_v1"

    /**
     * Distinct from [AAD_SLOT_A] so a recovery-escrow blob can never be swapped in
     * as a slot-A blob (GCM auth covers the AAD). Used by Android's whole-payload
     * offline-restore escrow.
     */
    const val AAD_RECOVERY = "recovery_escrow_v1"

    /**
     * AAD for the browser-byte-compatible root-key escrow. The browser wraps the
     * 32-byte root key under AAD = the single byte 0x01 (`vault.ts:298`); matching
     * it byte-for-byte is what lets the same 12 words unwrap the same root key on
     * either platform. Do NOT change — it is a cross-platform wire constant.
     */
    val AAD_RECOVERY_ROOTKEY = byteArrayOf(1)

    private val secureRandom = SecureRandom()
    private var _lastTimestamp = 0L

    /**
     * MF-4 / LV-1: NFC-normalize the password, then hash explicit UTF-8 bytes —
     * byte-for-byte identical to the browser's `password.normalize("NFC")` fed
     * to Argon2id (`argon2.ts` deriveKey). NFC normalization makes a visually
     * identical password typed in a different normalization form (NFD on macOS
     * dead-keys, some IMEs) derive the same key on every platform. UTF-8 (not
     * BouncyCastle's per-UTF-16-unit char[] path) keeps non-BMP characters
     * (emoji, surrogate pairs) consistent across platforms.
     */
    fun deriveKey(password: String, salt: ByteArray): ByteArray =
        deriveKeyFromBytes(
            Normalizer.normalize(password, Normalizer.Form.NFC).toByteArray(Charsets.UTF_8),
            salt,
        )

    /**
     * LV-1 legacy fallback ONLY: derive WITHOUT NFC normalization (raw UTF-8
     * bytes), for unlocking a vault sealed before normalization existed with a
     * non-NFC password. Every normal path uses [deriveKey].
     */
    fun deriveKeyRaw(password: String, salt: ByteArray): ByteArray =
        deriveKeyFromBytes(password.toByteArray(Charsets.UTF_8), salt)

    private fun deriveKeyFromBytes(passwordBytes: ByteArray, salt: ByteArray): ByteArray {
        val params = Argon2Parameters.Builder(Argon2Parameters.ARGON2_id)
            .withSalt(salt)
            .withParallelism(ARGON2_PARALLELISM)
            .withMemoryAsKB(ARGON2_MEMORY_KB)
            .withIterations(ARGON2_ITERATIONS)
            .build()
        val generator = Argon2BytesGenerator().apply { init(params) }
        val hash = ByteArray(ARGON2_HASH_LENGTH)
        try {
            generator.generateBytes(passwordBytes, hash)
        } finally {
            // Zero the transient password bytes immediately after use.
            passwordBytes.fill(0)
        }
        return hash
    }

    fun encrypt(plaintext: ByteArray, key: ByteArray, aad: String): ByteArray =
        encrypt(plaintext, key, aad.toByteArray(Charsets.UTF_8))

    fun decrypt(data: ByteArray, key: ByteArray, aad: String): ByteArray =
        decrypt(data, key, aad.toByteArray(Charsets.UTF_8))

    /**
     * Raw-AAD encrypt. Output is `nonce(12) || ciphertext || tag(16)`. The
     * [aad] bytes are authenticated, not encrypted. Used for the root-key
     * escrow whose AAD is the browser's single 0x01 byte ([AAD_RECOVERY_ROOTKEY]).
     */
    fun encrypt(plaintext: ByteArray, key: ByteArray, aad: ByteArray): ByteArray {
        val nonce = ByteArray(NONCE_SIZE).also(secureRandom::nextBytes)
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(
            Cipher.ENCRYPT_MODE,
            SecretKeySpec(key, "AES"),
            GCMParameterSpec(TAG_LENGTH_BITS, nonce),
        )
        cipher.updateAAD(aad)
        return nonce + cipher.doFinal(plaintext)
    }

    fun decrypt(data: ByteArray, key: ByteArray, aad: ByteArray): ByteArray {
        require(data.size >= NONCE_SIZE) { "Data too short" }
        val nonce = data.copyOfRange(0, NONCE_SIZE)
        val ciphertext = data.copyOfRange(NONCE_SIZE, data.size)
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(
            Cipher.DECRYPT_MODE,
            SecretKeySpec(key, "AES"),
            GCMParameterSpec(TAG_LENGTH_BITS, nonce),
        )
        cipher.updateAAD(aad)
        return cipher.doFinal(ciphertext)
    }

    fun generateSalt(): ByteArray = ByteArray(SALT_SIZE).also(secureRandom::nextBytes)
    fun generateRootKey(): ByteArray = ByteArray(KEY_SIZE).also(secureRandom::nextBytes)
    fun generateNoise(): ByteArray = ByteArray(DURESS_SLOT_SIZE).also(secureRandom::nextBytes)

    fun padToBucket(data: ByteArray): ByteArray {
        val target = PADDING_BUCKETS.firstOrNull { it >= data.size } ?: data.size
        return data.copyOf(target)
    }

    fun padToBlock(data: ByteArray): ByteArray {
        val remainder = data.size % PAD_BLOCK
        if (remainder == 0) return data
        return data.copyOf(data.size + (PAD_BLOCK - remainder))
    }

    @Synchronized
    fun monotonicEntryTimestamp(currentModifiedAt: Long? = null): Long {
        val base = maxOf(System.currentTimeMillis(), (currentModifiedAt ?: 0L) + 1L, _lastTimestamp + 1L)
        _lastTimestamp = base
        return base
    }

    /**
     * Reset the process-global monotonic entry clock. TEST ONLY.
     *
     * [monotonicEntryTimestamp] never moves backwards, so any test that feeds
     * it a far-future baseline leaves every later test in the same JVM issuing
     * far-future timestamps. That is invisible until something compares against
     * a clock-skew bound — [SyncLimits.clampSyncTimestamp] does, and a local
     * entry stamped beyond the skew window silently strands remote tombstones.
     */
    fun resetMonotonicClockForTest() {
        _lastTimestamp = 0L
    }

    fun zero(bytes: ByteArray) {
        bytes.fill(0)
    }
}