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)
    private var supported: Boolean? = null

    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()
        when (prefs.getInt(KEY_FORMAT, 0)) {
            FORMAT_VERSION -> Unit
            0 -> clearKnownKeystoreAliases()
            else -> clear()
        }
    }

    /**
     * Whether this device can host a PIN at all. Measured by probing the Keystore
     * rather than by reading FEATURE_KEYSTORE_LIMITED_USE_KEY / SINGLE_USE_KEY:
     * a Galaxy Z Fold 6 on SDK 36 enforces max-usage-1 HMAC keys in both
     * StrongBox and TEE while declaring neither feature, so the flags would
     * disable PIN unlock on hardware that supports it outright.
     */
    fun isSupported(): Boolean {
        supported?.let { return it }
        val measured = Build.VERSION.SDK_INT >= Build.VERSION_CODES.S && probeLimitedUseEnforcement()
        supported = measured
        return measured
    }

    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) &&
            existingSecretKey(VAULT_ALIAS) != null &&
            existingSecretKey(VERIFIER_ALIAS) != null

    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)
            val attemptAliases = 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))
                .putString(KEY_ATTEMPT_ALIASES, attemptAliases.joinToString(DELIMITER))
                .putBoolean(KEY_PREFERRED, false)
                .commit()
            if (!saved) throw java.io.IOException("pin_storage_commit_failed")
            true
        } catch (_: Exception) {
            createdAliases.forEach(::deleteKey)
            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
        var salt: ByteArray? = null
        var expected: ByteArray? = null
        var pinKey: ByteArray? = null
        var actual: ByteArray? = null
        var wrapped: ByteArray? = null
        var iv: ByteArray? = null
        return try {
            salt = prefs.getString(KEY_SALT, null)?.let(::decode) ?: return null
            expected = prefs.getString(KEY_VERIFIER, null)?.let(::decode) ?: return null
            val verifierKey = existingSecretKey(VERIFIER_ALIAS) ?: return null
            pinKey = VaultCrypto.deriveKey(pin, salt!!)
            actual = hmac(verifierKey, pinKey)
            if (!MessageDigest.isEqual(expected, actual)) return null

            wrapped = prefs.getString(KEY_WRAPPED_KEY, null)?.let(::decode) ?: return null
            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 {
            salt?.let(VaultCrypto::zero)
            expected?.let(VaultCrypto::zero)
            pinKey?.let(VaultCrypto::zero)
            actual?.let(VaultCrypto::zero)
            wrapped?.let(VaultCrypto::zero)
            iv?.let(VaultCrypto::zero)
        }
    }

    fun areAttemptsExhausted(): Boolean = isConfigured() &&
        loadAttemptAliases().none { existingSecretKey(it) != null }

    /**
     * Prove this device actually enforces a one-use hardware HMAC key instead of
     * trusting its feature flag: the first MAC must succeed and the same key must
     * then be unavailable or reject a second MAC. Generation also requires a
     * hardware-backed key, so a software-only provider fails the probe.
     */
    internal fun probeLimitedUseEnforcement(): Boolean {
        val alias = "$PROBE_ALIAS_PREFIX${UUID.randomUUID()}"
        val createdAliases = mutableListOf<String>()
        return try {
            val key = generateHmacKey(alias, limitedUse = true, createdAliases)
            hmac(key, ATTEMPT_CHALLENGE).let(VaultCrypto::zero)
            val remainingKey = existingSecretKey(alias) ?: return true
            val secondUseSucceeded = runCatching {
                hmac(remainingKey, ATTEMPT_CHALLENGE).let(VaultCrypto::zero)
            }.isSuccess
            !secondUseSucceeded
        } catch (_: Exception) {
            false
        } finally {
            deleteKey(alias)
        }
    }

    /** Regenerate the hardware attempt budget after a master-password unlock. */
    fun replenishAttemptGates() {
        if (!isConfigured()) return
        val oldAliases = loadAttemptAliases()
        val newAliases = mutableListOf<String>()
        try {
            val attemptAliases = createAttemptGates(newAliases)
            if (!prefs.edit().putString(KEY_ATTEMPT_ALIASES, attemptAliases.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>): List<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")
        return aliases
    }

    @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 clearKnownKeystoreAliases() {
        runCatching {
            KeyStore.getInstance(KEYSTORE).apply { load(null) }.aliases().toList()
                .filter {
                    it == VAULT_ALIAS || it == VERIFIER_ALIAS ||
                        it.startsWith(ATTEMPT_ALIAS_PREFIX) || it.startsWith(PROBE_ALIAS_PREFIX)
                }
                .forEach(::deleteKey)
        }
    }

    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 = 3
        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 PROBE_ALIAS_PREFIX = "peach_pin_probe_v3_"
        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)
    }
}