package com.peachpasswords.vault

import org.bouncycastle.crypto.generators.Argon2BytesGenerator
import org.bouncycastle.crypto.params.Argon2Parameters
import java.security.SecureRandom
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

    fun deriveKey(password: String, 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)
        generator.generateBytes(password.toCharArray(), hash)
        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
    }

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