package no.zyon.shade.keystore

import android.os.Build
import android.security.keystore.KeyGenParameterSpec
import android.security.keystore.KeyProperties
import javax.crypto.Cipher
import javax.crypto.KeyGenerator
import javax.crypto.SecretKey
import javax.crypto.spec.GCMParameterSpec
import java.security.KeyStore

/**
 * Hardware-backed AES-256-GCM master key in the Android Keystore.
 *
 * The key never leaves the secure environment — Android's keystore
 * implementation enforces that all encrypt/decrypt operations
 * happen inside the TEE (or StrongBox if present), and the raw
 * key bytes are never returned to userspace.
 *
 * The key is created on first use with these properties:
 *
 *  - AES-256-GCM, no padding
 *  - User authentication required: opt-in via the `requireBiometric`
 *    flag. When true, every encrypt/decrypt operation must be wrapped
 *    in a `BiometricPrompt.authenticate(CryptoObject(cipher))` call
 *    that succeeds within the same `Cipher` instance.
 *  - StrongBox-backed if available (Pixel 3+, most Samsung flagships).
 *    Falls back to TEE on devices without StrongBox.
 *  - InvalidatedByBiometricEnrollment(true): a newly enrolled
 *    fingerprint/face invalidates the key, forcing the user to
 *    re-bootstrap from credentials. Defends against a thief who
 *    enrolls their own biometric.
 *
 * Mirrors the role `KeyManager` plays in `@shade/storage-encrypted`'s
 * V4.5 KDF chain: this is the *encryption-at-rest* master key, not
 * the X3DH identity key. The Shade protocol's identity keys are
 * stored encrypted under THIS key.
 */
class KeystoreMasterKey(
    private val alias: String,
    private val requireBiometric: Boolean = true,
) {

    init {
        require(alias.isNotEmpty()) { "alias must be non-empty" }
    }

    /**
     * Build a `Cipher` initialized for encryption with the master key.
     *
     * If the key requires user auth, the returned cipher is *not yet
     * usable* — the caller MUST wrap it in a
     * `BiometricPrompt.authenticate(CryptoObject(cipher))` and use
     * the cipher exposed by the auth-success callback. Calling
     * `cipher.doFinal(...)` before authentication throws
     * `UserNotAuthenticatedException`.
     */
    fun cipherForEncrypt(): Cipher {
        val key = getOrCreateKey()
        val cipher = Cipher.getInstance(TRANSFORMATION)
        cipher.init(Cipher.ENCRYPT_MODE, key)
        return cipher
    }

    /**
     * Build a `Cipher` initialized for decryption with the master key
     * and a previously-stored 12-byte nonce. Same authentication
     * requirement as `cipherForEncrypt`.
     */
    fun cipherForDecrypt(nonce: ByteArray): Cipher {
        require(nonce.size == 12) { "GCM nonce must be 12 bytes" }
        val key = getOrCreateKey()
        val cipher = Cipher.getInstance(TRANSFORMATION)
        cipher.init(Cipher.DECRYPT_MODE, key, GCMParameterSpec(128, nonce))
        return cipher
    }

    /** True if the key already exists in the Android Keystore. */
    fun exists(): Boolean {
        val ks = KeyStore.getInstance(ANDROID_KEYSTORE).apply { load(null) }
        return ks.containsAlias(alias)
    }

    /**
     * Delete the master key. Catastrophic — all data encrypted under
     * it becomes unrecoverable. Used by the "forget everything" flow
     * (paired with `Profile.delete()` in the V4.9 namespace).
     */
    fun deleteKey() {
        val ks = KeyStore.getInstance(ANDROID_KEYSTORE).apply { load(null) }
        if (ks.containsAlias(alias)) ks.deleteEntry(alias)
    }

    private fun getOrCreateKey(): SecretKey {
        val ks = KeyStore.getInstance(ANDROID_KEYSTORE).apply { load(null) }
        ks.getEntry(alias, null)?.let { entry ->
            return (entry as KeyStore.SecretKeyEntry).secretKey
        }
        return generateKey()
    }

    private fun generateKey(): 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)
            // Each encrypt operation generates a fresh IV in the secure
            // env; we read it back via `cipher.iv` after init.
            .setRandomizedEncryptionRequired(true)

        if (requireBiometric) {
            builder.setUserAuthenticationRequired(true)
            // BIOMETRIC_STRONG only — class 3, the highest assurance
            // level (Class 3 = false-accept rate < 1/50 000 per BiometricPrompt).
            // DEVICE_CREDENTIAL is intentionally NOT included: a stolen
            // device with a known PIN should not unlock Shade.
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
                builder.setUserAuthenticationParameters(
                    /* timeout = */ 0,
                    KeyProperties.AUTH_BIOMETRIC_STRONG,
                )
            } else {
                @Suppress("DEPRECATION")
                builder.setUserAuthenticationValidityDurationSeconds(-1)
            }
            builder.setInvalidatedByBiometricEnrollment(true)
        }

        // StrongBox if available — bumps key storage to a dedicated
        // tamper-resistant chip on Pixel 3+ / most Samsung flagships.
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.P) {
            builder.setIsStrongBoxBacked(true)
        }

        val gen = KeyGenerator.getInstance(KeyProperties.KEY_ALGORITHM_AES, ANDROID_KEYSTORE)
        return try {
            gen.init(builder.build())
            gen.generateKey()
        } catch (_: Exception) {
            // StrongBox not present or full → retry without StrongBox.
            // Same for older devices that don't honor
            // setUserAuthenticationParameters.
            val fallback = KeyGenParameterSpec.Builder(
                alias,
                KeyProperties.PURPOSE_ENCRYPT or KeyProperties.PURPOSE_DECRYPT,
            )
                .setBlockModes(KeyProperties.BLOCK_MODE_GCM)
                .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_NONE)
                .setKeySize(256)
                .setRandomizedEncryptionRequired(true)
                .apply {
                    if (requireBiometric) {
                        setUserAuthenticationRequired(true)
                        @Suppress("DEPRECATION")
                        setUserAuthenticationValidityDurationSeconds(-1)
                        setInvalidatedByBiometricEnrollment(true)
                    }
                }
                .build()
            gen.init(fallback)
            gen.generateKey()
        }
    }

    companion object {
        private const val ANDROID_KEYSTORE = "AndroidKeyStore"
        private const val TRANSFORMATION = "AES/GCM/NoPadding"
    }
}
