package com.peachpasswords.vault

import android.content.Context
import android.content.pm.PackageManager
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 dagger.hilt.android.qualifiers.ApplicationContext
import java.security.KeyStore
import java.security.MessageDigest
import java.util.Base64
import java.util.UUID
import javax.crypto.Cipher
import javax.crypto.KeyGenerator
import javax.crypto.Mac
import javax.crypto.SecretKey
import javax.crypto.SecretKeyFactory
import javax.crypto.spec.GCMParameterSpec
import javax.inject.Inject
import javax.inject.Singleton

internal const val PIN_MIN_LENGTH = 4
internal const val PIN_MAX_LENGTH = 12

internal fun pinValidationError(pin: String): String? = when {
    pin.length !in PIN_MIN_LENGTH..PIN_MAX_LENGTH -> "PIN must be 4–12 digits"
    pin.any { it !in '0'..'9' } -> "PIN must contain digits only"
    else -> null
}

/**
 * Device-local PIN gate for a vault key held by Android Keystore.
 *
 * The PIN is never used to encrypt the vault key. A hardware-backed AES key
 * encrypts the vault key directly, while a separate hardware-backed HMAC key
 * verifies Argon2id(PIN). Every attempt first consumes one single-use HMAC gate
 * in Keystore. The salt, verifier, IV, and ciphertext in SharedPreferences are
 * therefore useless for offline guessing without the device's Keystore keys.
 */
@Singleton
class PinUnlockStore @Inject constructor(
    @ApplicationContext context: Context,
) {
    private val appContext = context.applicationContext
    private val prefs = appContext.getSharedPreferences(PREFS, Context.MODE_PRIVATE)

    init {
        // The v1 wrapper was directly brute-forceable from a copied prefs file.
        // Remove it as soon as the security-remediated app opens the store; users
        // can still unlock with the master password and configure the new format.
        clearLegacyConfiguration()
    }

    /** New PIN setup is safe only where Android reports hardware limited-use keys. */
    fun isSupported(): Boolean = Build.VERSION.SDK_INT >= Build.VERSION_CODES.S &&
        (appContext.packageManager.hasSystemFeature(PackageManager.FEATURE_KEYSTORE_LIMITED_USE_KEY) ||
            appContext.packageManager.hasSystemFeature(PackageManager.FEATURE_KEYSTORE_SINGLE_USE_KEY))

    fun isConfigured(): Boolean =
        prefs.getInt(KEY_FORMAT, 0) == FORMAT_VERSION &&
            prefs.contains(KEY_SALT) &&
            prefs.contains(KEY_VERIFIER) &&
            prefs.contains(KEY_WRAPPED_KEY) &&
            prefs.contains(KEY_WRAPPED_IV)

    fun isPinPreferred(): Boolean = prefs.getBoolean(KEY_PREFERRED, false)

    fun markPinUsed() {
        prefs.edit().putBoolean(KEY_PREFERRED, true).apply()
    }

    fun markPasswordUsed() {
        prefs.edit().putBoolean(KEY_PREFERRED, false).apply()
    }

    fun configure(pin: String, vaultKey: ByteArray): Boolean {
        if (!isSupported() || pinValidationError(pin) != null || vaultKey.size != VaultCrypto.KEY_SIZE) return false
        clear()

        var salt: ByteArray? = null
        var pinKey: ByteArray? = null
        var verifier: ByteArray? = null
        var wrapped: ByteArray? = null
        val createdAliases = mutableListOf<String>()
        return try {
            val vaultKeystoreKey = generateAesKey(VAULT_ALIAS, createdAliases)
            val verifierKeystoreKey = generateHmacKey(VERIFIER_ALIAS, limitedUse = false, createdAliases)
            salt = VaultCrypto.generateSalt()
            pinKey = VaultCrypto.deriveKey(pin, salt)
            verifier = hmac(verifierKeystoreKey, pinKey)

            val cipher = Cipher.getInstance(TRANSFORM).apply { init(Cipher.ENCRYPT_MODE, vaultKeystoreKey) }
            wrapped = cipher.doFinal(vaultKey)
            createAttemptGates(createdAliases)

            val saved = prefs.edit()
                .putInt(KEY_FORMAT, FORMAT_VERSION)
                .putString(KEY_SALT, b64(salt!!))
                .putString(KEY_VERIFIER, b64(verifier!!))
                .putString(KEY_WRAPPED_KEY, b64(wrapped!!))
                .putString(KEY_WRAPPED_IV, b64(cipher.iv))
                .putBoolean(KEY_PREFERRED, false)
                .commit()
            if (!saved) throw java.io.IOException("pin_storage_commit_failed")
            true
        } catch (_: Exception) {
            clear()
            false
        } finally {
            salt?.let(VaultCrypto::zero)
            pinKey?.let(VaultCrypto::zero)
            verifier?.let(VaultCrypto::zero)
            wrapped?.let(VaultCrypto::zero)
        }
    }

    /**
     * Returns the vault key only after a hardware gate has been consumed and the
     * hardware HMAC verifier accepts the Argon2id-derived PIN key.
     */
    fun unlockVaultKey(pin: String): ByteArray? {
        if (pinValidationError(pin) != null || !isConfigured() || !consumeAttemptGate()) return null
        val salt = prefs.getString(KEY_SALT, null)?.let(::decode) ?: return null
        val expected = prefs.getString(KEY_VERIFIER, null)?.let(::decode) ?: return null
        var pinKey: ByteArray? = null
        var actual: ByteArray? = null
        return try {
            val verifierKey = existingSecretKey(VERIFIER_ALIAS) ?: return null
            pinKey = VaultCrypto.deriveKey(pin, salt)
            actual = hmac(verifierKey, pinKey)
            if (!MessageDigest.isEqual(expected, actual)) return null

            val wrapped = prefs.getString(KEY_WRAPPED_KEY, null)?.let(::decode) ?: return null
            val iv = prefs.getString(KEY_WRAPPED_IV, null)?.let(::decode) ?: return null
            val vaultKey = Cipher.getInstance(TRANSFORM).run {
                init(Cipher.DECRYPT_MODE, existingSecretKey(VAULT_ALIAS) ?: return null, GCMParameterSpec(TAG_BITS, iv))
                doFinal(wrapped)
            }
            // A successful PIN proves possession of the PIN; replenish the
            // hardware attempt budget for the next lock/unlock cycle.
            replenishAttemptGates()
            vaultKey
        } catch (_: Exception) {
            null
        } finally {
            VaultCrypto.zero(salt)
            VaultCrypto.zero(expected)
            pinKey?.let(VaultCrypto::zero)
            actual?.let(VaultCrypto::zero)
        }
    }

    fun areAttemptsExhausted(): Boolean = isConfigured() && loadAttemptAliases().isEmpty()

    /** Regenerate the hardware attempt budget after a master-password unlock. */
    fun replenishAttemptGates() {
        if (!isConfigured()) return
        val oldAliases = loadAttemptAliases()
        val newAliases = mutableListOf<String>()
        try {
            createAttemptGates(newAliases)
            if (!prefs.edit().putString(KEY_ATTEMPT_ALIASES, newAliases.joinToString(DELIMITER)).commit()) {
                throw java.io.IOException("pin_attempt_gate_commit_failed")
            }
            oldAliases.forEach(::deleteKey)
        } catch (_: Exception) {
            newAliases.forEach(::deleteKey)
        }
    }

    fun clear() {
        loadAttemptAliases().forEach(::deleteKey)
        listOf(VAULT_ALIAS, VERIFIER_ALIAS).forEach(::deleteKey)
        clearLegacyConfiguration()
        prefs.edit()
            .remove(KEY_FORMAT)
            .remove(KEY_SALT)
            .remove(KEY_VERIFIER)
            .remove(KEY_WRAPPED_KEY)
            .remove(KEY_WRAPPED_IV)
            .remove(KEY_ATTEMPT_ALIASES)
            .remove(KEY_PREFERRED)
            .commit()
    }

    private fun consumeAttemptGate(): Boolean {
        val aliases = loadAttemptAliases()
        for (alias in aliases) {
            val key = existingSecretKey(alias) ?: continue
            try {
                // The output is intentionally discarded; the operation exists to
                // consume the one-use hardware authorization before KDF work.
                hmac(key, ATTEMPT_CHALLENGE)
                deleteKey(alias)
                prefs.edit().putString(
                    KEY_ATTEMPT_ALIASES,
                    aliases.filterNot { it == alias }.joinToString(DELIMITER),
                ).commit()
                return true
            } catch (_: KeyPermanentlyInvalidatedException) {
                deleteKey(alias)
            } catch (_: Exception) {
                deleteKey(alias)
            }
        }
        return false
    }

    private fun createAttemptGates(createdAliases: MutableList<String>) {
        repeat(ATTEMPT_GATE_COUNT) {
            val alias = "$ATTEMPT_ALIAS_PREFIX${UUID.randomUUID()}"
            generateHmacKey(alias, limitedUse = true, createdAliases)
        }
        val aliases = createdAliases.filter { it.startsWith(ATTEMPT_ALIAS_PREFIX) }
        if (aliases.size != ATTEMPT_GATE_COUNT) throw java.io.IOException("pin_attempt_gate_count")
        prefs.edit().putString(KEY_ATTEMPT_ALIASES, aliases.joinToString(DELIMITER)).commit()
    }

    @Suppress("NewApi")
    private fun generateAesKey(alias: String, createdAliases: MutableList<String>): SecretKey {
        existingSecretKey(alias)?.let { key ->
            if (isHardwareBacked(key)) return key
            deleteKey(alias)
        }
        val strongBoxPreferred = Build.VERSION.SDK_INT >= Build.VERSION_CODES.P &&
            appContext.packageManager.hasSystemFeature(PackageManager.FEATURE_STRONGBOX_KEYSTORE)
        val key = try {
            generateAesKey(alias, strongBoxPreferred)
        } catch (_: StrongBoxUnavailableException) {
            deleteKey(alias)
            generateAesKey(alias, false)
        }
        if (!isHardwareBacked(key)) {
            deleteKey(alias)
            throw java.security.GeneralSecurityException("pin_key_not_hardware_backed")
        }
        createdAliases += alias
        return key
    }

    @Suppress("NewApi")
    private fun generateHmacKey(alias: String, limitedUse: Boolean, createdAliases: MutableList<String>): SecretKey {
        val strongBoxPreferred = Build.VERSION.SDK_INT >= Build.VERSION_CODES.P &&
            appContext.packageManager.hasSystemFeature(PackageManager.FEATURE_STRONGBOX_KEYSTORE)
        val key = try {
            generateHmacKey(alias, limitedUse, strongBoxPreferred)
        } catch (_: StrongBoxUnavailableException) {
            deleteKey(alias)
            generateHmacKey(alias, limitedUse, false)
        }
        if (!isHardwareBacked(key)) {
            deleteKey(alias)
            throw java.security.GeneralSecurityException("pin_key_not_hardware_backed")
        }
        createdAliases += alias
        return key
    }

    private fun generateAesKey(alias: String, strongBox: Boolean): SecretKey {
        val builder = KeyGenParameterSpec.Builder(
            alias,
            KeyProperties.PURPOSE_ENCRYPT or KeyProperties.PURPOSE_DECRYPT,
        )
            .setBlockModes(KeyProperties.BLOCK_MODE_GCM)
            .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_NONE)
            .setKeySize(256)
        if (strongBox && Build.VERSION.SDK_INT >= Build.VERSION_CODES.P) builder.setIsStrongBoxBacked(true)
        val generator = KeyGenerator.getInstance(KeyProperties.KEY_ALGORITHM_AES, KEYSTORE)
        generator.init(builder.build())
        return generator.generateKey()
    }

    private fun generateHmacKey(alias: String, limitedUse: Boolean, strongBox: Boolean): SecretKey {
        val builder = KeyGenParameterSpec.Builder(alias, KeyProperties.PURPOSE_SIGN)
            .setDigests(KeyProperties.DIGEST_SHA256)
            .setKeySize(256)
        if (limitedUse && Build.VERSION.SDK_INT >= Build.VERSION_CODES.S) builder.setMaxUsageCount(1)
        if (strongBox && Build.VERSION.SDK_INT >= Build.VERSION_CODES.P) builder.setIsStrongBoxBacked(true)
        val generator = KeyGenerator.getInstance(KeyProperties.KEY_ALGORITHM_HMAC_SHA256, KEYSTORE)
        generator.init(builder.build())
        return generator.generateKey()
    }

    private fun hmac(key: SecretKey, input: ByteArray): ByteArray =
        Mac.getInstance("HmacSHA256").apply { init(key) }.doFinal(input)

    private fun existingSecretKey(alias: String): SecretKey? = runCatching {
        val keyStore = KeyStore.getInstance(KEYSTORE).apply { load(null) }
        (keyStore.getEntry(alias, null) as? KeyStore.SecretKeyEntry)?.secretKey
    }.getOrNull()

    private fun isHardwareBacked(key: SecretKey): Boolean = runCatching {
        val factory = SecretKeyFactory.getInstance(key.algorithm, KEYSTORE)
        val info = factory.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 loadAttemptAliases(): List<String> =
        prefs.getString(KEY_ATTEMPT_ALIASES, null)
            ?.split(DELIMITER)
            ?.filter { it.startsWith(ATTEMPT_ALIAS_PREFIX) }
            ?: emptyList()

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

    private fun clearLegacyConfiguration() {
        prefs.edit()
            .remove("peach_pin_salt")
            .remove("peach_pin_verifier")
            .remove("peach_pin_wrapped_key")
            .remove("peach_pin_preferred")
            .commit()
    }

    private fun b64(bytes: ByteArray): String = Base64.getEncoder().withoutPadding().encodeToString(bytes)
    private fun decode(value: String): ByteArray = Base64.getDecoder().decode(value)

    private companion object {
        const val FORMAT_VERSION = 2
        const val PREFS = "peach_vault_store"
        const val KEY_FORMAT = "peach_pin_format"
        const val KEY_SALT = "peach_pin_salt_v2"
        const val KEY_VERIFIER = "peach_pin_verifier_v2"
        const val KEY_WRAPPED_KEY = "peach_pin_wrapped_key_v2"
        const val KEY_WRAPPED_IV = "peach_pin_wrapped_iv_v2"
        const val KEY_ATTEMPT_ALIASES = "peach_pin_attempt_aliases_v2"
        const val KEY_PREFERRED = "peach_pin_preferred_v2"
        const val KEYSTORE = "AndroidKeyStore"
        const val VAULT_ALIAS = "peach_pin_vault_key_v2"
        const val VERIFIER_ALIAS = "peach_pin_verifier_key_v2"
        const val ATTEMPT_ALIAS_PREFIX = "peach_pin_attempt_v2_"
        const val ATTEMPT_GATE_COUNT = 5
        const val TRANSFORM = "AES/GCM/NoPadding"
        const val TAG_BITS = 128
        const val DELIMITER = ","
        val ATTEMPT_CHALLENGE = "peach_pin_attempt_gate_v2".toByteArray(Charsets.UTF_8)
    }
}