package com.peachpasswords.ui.lock

import android.content.Context
import android.os.Build
import android.security.keystore.KeyGenParameterSpec
import android.security.keystore.KeyInfo
import android.security.keystore.KeyPermanentlyInvalidatedException
import android.security.keystore.KeyProperties
import android.security.keystore.StrongBoxUnavailableException
import java.security.KeyStore
import java.security.GeneralSecurityException
import java.util.Base64
import javax.crypto.Cipher
import javax.crypto.KeyGenerator
import javax.crypto.SecretKey
import javax.crypto.SecretKeyFactory
import javax.crypto.spec.GCMParameterSpec

/**
 * Stores the vault's AES key wrapped under a hardware-backed, biometric-gated
 * Android KeyStore key, so the lock screen can unlock without the master password.
 *
 * The KeyStore key requires user (biometric) authentication for every use, so both
 * wrapping (on enable) and unwrapping (on unlock) must happen inside a
 * [androidx.biometric.BiometricPrompt] `CryptoObject`. The wrapped blob + IV are
 * kept in plain prefs — they are AES-256-GCM ciphertext under a key that never
 * leaves the secure hardware, so the prefs themselves carry no secret.
 *
 * Enrolling a new biometric invalidates the key ([setInvalidatedByBiometricEnrollment]);
 * the resulting [KeyPermanentlyInvalidatedException] is caught and the escrow cleared,
 * so the user simply falls back to the master password.
 */
object BiometricKeyStore {
    private const val KEY_ALIAS = "peach_bio_vault_key"
    private const val KEYSTORE = "AndroidKeyStore"
    private const val TRANSFORM = "AES/GCM/NoPadding"
    private const val PREFS = "peach_bio_prefs"
    private const val K_WRAPPED = "wrapped"
    private const val K_IV = "iv"
    private const val GCM_TAG_BITS = 128

    private fun prefs(context: Context) =
        context.getSharedPreferences(PREFS, Context.MODE_PRIVATE)

    fun hasWrappedKey(context: Context): Boolean =
        prefs(context).contains(K_WRAPPED) && prefs(context).contains(K_IV) &&
            (runCatching { existingKey()?.let(::isHardwareBacked) == true }.getOrDefault(false))

    fun clear(context: Context) {
        prefs(context).edit().remove(K_WRAPPED).remove(K_IV).apply()
        runCatching {
            KeyStore.getInstance(KEYSTORE).apply { load(null) }.deleteEntry(KEY_ALIAS)
        }
    }

    /** Cipher to wrap the vault key on enable (fresh random IV). Drive via CryptoObject. */
    fun encryptCipher(): Cipher {
        val key = getOrCreateKey()
        return Cipher.getInstance(TRANSFORM).apply { init(Cipher.ENCRYPT_MODE, key) }
    }

    /**
     * Cipher to unwrap on unlock, initialised with the stored IV. Returns null when
     * no escrow exists or the key was invalidated by a new biometric enrolment (the
     * escrow is cleared in that case so the caller falls back to the password).
     */
    fun decryptCipher(context: Context): Cipher? {
        val iv = loadIv(context) ?: return null
        return try {
            val key = existingKey()?.takeIf(::isHardwareBacked) ?: return null
            Cipher.getInstance(TRANSFORM).apply { init(Cipher.DECRYPT_MODE, key, GCMParameterSpec(GCM_TAG_BITS, iv)) }
        } catch (_: KeyPermanentlyInvalidatedException) {
            clear(context)
            null
        } catch (_: Exception) {
            null
        }
    }

    /** Wrap [vaultKey] with the authenticated [cipher] and persist it. */
    fun wrapAndStore(context: Context, cipher: Cipher, vaultKey: ByteArray) {
        val wrapped = cipher.doFinal(vaultKey)
        prefs(context).edit()
            .putString(K_WRAPPED, b64(wrapped))
            .putString(K_IV, b64(cipher.iv))
            .apply()
    }

    /** Unwrap the stored vault key with the authenticated [cipher]. */
    fun unwrap(context: Context, cipher: Cipher): ByteArray? {
        val wrapped = prefs(context).getString(K_WRAPPED, null)?.let(::unb64) ?: return null
        return runCatching { cipher.doFinal(wrapped) }.getOrNull()
    }

    private fun loadIv(context: Context): ByteArray? =
        prefs(context).getString(K_IV, null)?.let(::unb64)

    @Suppress("NewApi")
    private fun getOrCreateKey(): SecretKey {
        existingKey()?.let { key ->
            if (isHardwareBacked(key)) return key
            deleteKey()
        }
        val builder = KeyGenParameterSpec.Builder(
            KEY_ALIAS,
            KeyProperties.PURPOSE_ENCRYPT or KeyProperties.PURPOSE_DECRYPT,
        )
            .setBlockModes(KeyProperties.BLOCK_MODE_GCM)
            .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_NONE)
            .setKeySize(256)
            .setUserAuthenticationRequired(true)
            .setInvalidatedByBiometricEnrollment(true)
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
            builder.setUserAuthenticationParameters(0, KeyProperties.AUTH_BIOMETRIC_STRONG)
        }
        val key = try {
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.P) {
                builder.setIsStrongBoxBacked(true)
            }
            generateKey(builder)
        } catch (_: StrongBoxUnavailableException) {
            generateKey(builder, strongBox = false)
        }
        if (!isHardwareBacked(key)) {
            deleteKey()
            throw GeneralSecurityException("biometric_key_not_hardware_backed")
        }
        return key
    }

    private fun generateKey(builder: KeyGenParameterSpec.Builder, strongBox: Boolean = true): SecretKey {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.P && !strongBox) {
            builder.setIsStrongBoxBacked(false)
        }
        val gen = KeyGenerator.getInstance(KeyProperties.KEY_ALGORITHM_AES, KEYSTORE)
        gen.init(builder.build())
        return gen.generateKey()
    }

    private fun existingKey(): SecretKey? {
        val ks = KeyStore.getInstance(KEYSTORE).apply { load(null) }
        return (ks.getEntry(KEY_ALIAS, null) as? KeyStore.SecretKeyEntry)?.secretKey
    }

    private fun isHardwareBacked(key: SecretKey): Boolean = runCatching {
        val info = SecretKeyFactory.getInstance(key.algorithm, KEYSTORE)
            .getKeySpec(key, KeyInfo::class.java) as KeyInfo
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.S) {
            info.securityLevel == KeyProperties.SECURITY_LEVEL_TRUSTED_ENVIRONMENT ||
                info.securityLevel == KeyProperties.SECURITY_LEVEL_STRONGBOX
        } else {
            @Suppress("DEPRECATION")
            info.isInsideSecureHardware
        }
    }.getOrDefault(false)

    private fun deleteKey() {
        runCatching { KeyStore.getInstance(KEYSTORE).apply { load(null) }.deleteEntry(KEY_ALIAS) }
    }

    private fun b64(b: ByteArray) = Base64.getEncoder().encodeToString(b)
    private fun unb64(s: String) = Base64.getDecoder().decode(s)
}