package no.zyon.shade.keystore

import android.content.Context
import android.content.SharedPreferences
import android.util.Base64
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import no.zyon.shade.crypto.CryptoProvider
import no.zyon.shade.serialization.SessionStateJson
import no.zyon.shade.storage.StorageProvider
import no.zyon.shade.types.IdentityKeyPair
import no.zyon.shade.types.OneTimePreKey
import no.zyon.shade.types.SessionState
import no.zyon.shade.types.SignedPreKey

/**
 * `StorageProvider` implementation that gates all reads/writes through
 * a biometric-locked `KeystoreMasterKey`. Mirrors `MemoryStorage` for
 * the API surface but persists state to `SharedPreferences` with each
 * row encrypted under the keystore key.
 *
 * Lifecycle:
 *
 *   1. App start → construct `KeystoreStorage(context, alias)`.
 *   2. `unlock(BiometricUnlock)` runs the system biometric prompt.
 *      The Android keystore caches the auth state under the key's
 *      `setUserAuthenticationParameters(0, BIOMETRIC_STRONG)` policy
 *      until the next biometric event (re-enrollment, etc.).
 *   3. While unlocked, `getSession`/`saveSession` etc. work normally.
 *   4. `lock()` clears the in-memory unlocked flag so a future
 *      operation triggers another biometric prompt.
 *
 * Wire layout per row:
 *   `<base64(nonce(12))>:<base64(ct||tag)>`
 *
 * Stored as `String` SharedPreferences entries. AAD = the row's
 * preference key (`session:<address>`, `signedPreKey:<id>`, etc.) so
 * a substituted-prefs swap fails to open.
 */
class KeystoreStorage(
    context: Context,
    private val crypto: CryptoProvider,
    keyAlias: String = DEFAULT_KEY_ALIAS,
    prefsName: String = DEFAULT_PREFS_NAME,
    requireBiometric: Boolean = true,
) : StorageProvider {

    private val prefs: SharedPreferences =
        context.applicationContext.getSharedPreferences(prefsName, Context.MODE_PRIVATE)
    private val masterKey = KeystoreMasterKey(keyAlias, requireBiometric = requireBiometric)
    private val writeMutex = Mutex()

    @Volatile
    private var unlocked: Boolean = false

    /**
     * Unlock the keystore via biometric prompt. Idempotent — calling
     * twice without a `lock()` between is a no-op.
     */
    suspend fun unlock(unlock: BiometricUnlock) {
        if (unlocked) return
        // The biometric flow returns an authenticated *encrypt*
        // cipher; we discard it after a one-shot probe to confirm
        // the master key is reachable. The actual encrypt/decrypt
        // ciphers in the I/O path use the authentication state
        // established here (Android Keystore caches the auth for
        // the user-authentication-required key under the
        // `setUserAuthenticationParameters(0, BIOMETRIC_STRONG)`
        // policy until the next biometric event).
        val probe = masterKey.cipherForEncrypt()
        unlock.unlock(probe)
        unlocked = true
    }

    /** Unlock without biometric — only valid for keys constructed with `requireBiometric=false`. */
    fun unlockNoBiometric() {
        unlocked = true
    }

    /** Wipe in-memory unlock state. The key itself stays in the keystore. */
    fun lock() {
        unlocked = false
    }

    /**
     * Catastrophic reset: deletes the master key + all encrypted
     * preferences. Used by the "forget everything" / 3-strikes-wipe
     * path. The next bootstrap rebuilds from credentials.
     */
    fun forgetEverything() {
        masterKey.deleteKey()
        prefs.edit().clear().apply()
        unlocked = false
    }

    // ─── Identity ──────────────────────────────────────────────

    override suspend fun getIdentityKeyPair(): IdentityKeyPair? {
        val json = readDecrypted(KEY_IDENTITY) ?: return null
        return SessionStateJson.deserializeIdentityKeyPair(json)
    }

    override suspend fun saveIdentityKeyPair(keyPair: IdentityKeyPair) {
        writeEncrypted(KEY_IDENTITY, SessionStateJson.serializeIdentityKeyPair(keyPair))
    }

    override suspend fun getLocalRegistrationId(): Int {
        return readDecrypted(KEY_REGISTRATION_ID)?.toIntOrNull() ?: 0
    }

    override suspend fun saveLocalRegistrationId(id: Int) {
        writeEncrypted(KEY_REGISTRATION_ID, id.toString())
    }

    // ─── Signed prekeys ────────────────────────────────────────

    override suspend fun getSignedPreKey(keyId: Int): SignedPreKey? {
        val json = readDecrypted("$KEY_SIGNED_PREKEY:$keyId") ?: return null
        return SessionStateJson.deserializeSignedPreKey(json)
    }

    override suspend fun saveSignedPreKey(key: SignedPreKey) {
        writeEncrypted(
            "$KEY_SIGNED_PREKEY:${key.keyId}",
            SessionStateJson.serializeSignedPreKey(key),
        )
    }

    override suspend fun removeSignedPreKey(keyId: Int) {
        writeMutex.withLock { prefs.edit().remove("$KEY_SIGNED_PREKEY:$keyId").apply() }
    }

    // ─── One-time prekeys ──────────────────────────────────────

    override suspend fun getOneTimePreKey(keyId: Int): OneTimePreKey? {
        val json = readDecrypted("$KEY_ONETIME_PREKEY:$keyId") ?: return null
        return SessionStateJson.deserializeOneTimePreKey(json)
    }

    override suspend fun saveOneTimePreKey(key: OneTimePreKey) {
        writeEncrypted(
            "$KEY_ONETIME_PREKEY:${key.keyId}",
            SessionStateJson.serializeOneTimePreKey(key),
        )
    }

    override suspend fun removeOneTimePreKey(keyId: Int) {
        writeMutex.withLock { prefs.edit().remove("$KEY_ONETIME_PREKEY:$keyId").apply() }
    }

    override suspend fun getOneTimePreKeyCount(): Int {
        return prefs.all.keys.count { it.startsWith("$KEY_ONETIME_PREKEY:") }
    }

    // ─── Sessions ──────────────────────────────────────────────

    override suspend fun getSession(address: String): SessionState? {
        val json = readDecrypted("$KEY_SESSION:$address") ?: return null
        return SessionStateJson.deserialize(json)
    }

    override suspend fun saveSession(address: String, state: SessionState) {
        writeEncrypted("$KEY_SESSION:$address", SessionStateJson.serialize(state))
    }

    override suspend fun removeSession(address: String) {
        writeMutex.withLock { prefs.edit().remove("$KEY_SESSION:$address").apply() }
    }

    // ─── Trust ─────────────────────────────────────────────────

    override suspend fun isTrustedIdentity(address: String, identityKey: ByteArray): Boolean {
        val stored = readDecrypted("$KEY_TRUSTED:$address") ?: return true // TOFU
        val storedBytes = Base64.decode(stored, Base64.NO_WRAP)
        return crypto.constantTimeEqual(storedBytes, identityKey)
    }

    override suspend fun saveTrustedIdentity(address: String, identityKey: ByteArray) {
        writeEncrypted(
            "$KEY_TRUSTED:$address",
            Base64.encodeToString(identityKey, Base64.NO_WRAP),
        )
    }

    // ─── Encrypted-row plumbing ────────────────────────────────

    private fun ensureUnlocked() {
        check(unlocked) {
            "KeystoreStorage is locked — call unlock(BiometricUnlock) first"
        }
    }

    private suspend fun readDecrypted(prefKey: String): String? {
        ensureUnlocked()
        val raw = prefs.getString(prefKey, null) ?: return null
        val parts = raw.split(":", limit = 2)
        require(parts.size == 2) { "malformed encrypted row at $prefKey" }
        val nonce = Base64.decode(parts[0], Base64.NO_WRAP)
        val ct = Base64.decode(parts[1], Base64.NO_WRAP)
        val cipher = masterKey.cipherForDecrypt(nonce)
        cipher.updateAAD(prefKey.toByteArray(Charsets.UTF_8))
        return cipher.doFinal(ct).toString(Charsets.UTF_8)
    }

    private suspend fun writeEncrypted(prefKey: String, plaintext: String) {
        ensureUnlocked()
        writeMutex.withLock {
            val cipher = masterKey.cipherForEncrypt()
            cipher.updateAAD(prefKey.toByteArray(Charsets.UTF_8))
            val ct = cipher.doFinal(plaintext.toByteArray(Charsets.UTF_8))
            val nonce = cipher.iv
            val nonceB64 = Base64.encodeToString(nonce, Base64.NO_WRAP)
            val ctB64 = Base64.encodeToString(ct, Base64.NO_WRAP)
            prefs.edit().putString(prefKey, "$nonceB64:$ctB64").apply()
        }
    }

    companion object {
        const val DEFAULT_KEY_ALIAS = "shade-master-v1"
        const val DEFAULT_PREFS_NAME = "shade-keystore-storage-v1"

        private const val KEY_IDENTITY = "identity"
        private const val KEY_REGISTRATION_ID = "registrationId"
        private const val KEY_SIGNED_PREKEY = "signedPreKey"
        private const val KEY_ONETIME_PREKEY = "oneTimePreKey"
        private const val KEY_SESSION = "session"
        private const val KEY_TRUSTED = "trusted"
    }
}
