package app.closer.crypto

import app.closer.domain.model.Couple
import com.google.crypto.tink.Aead
import com.google.crypto.tink.KeysetHandle
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import javax.inject.Inject
import javax.inject.Singleton

enum class EncryptionStatus {
    /** Local keyset present -- ready to encrypt/decrypt. */
    UNLOCKED,
    /** Found keyset in the invite slot; moved to coupleId slot automatically (inviter's first load). */
    RECONCILED_FROM_INVITE,
    /** No local keyset on this device -- prompt for the recovery phrase. */
    NEEDS_RECOVERY
}

class MissingCoupleKeyException(coupleId: String) :
    IllegalStateException("Encrypted couple key is unavailable for $coupleId")

data class SetupResult(
    val handle: KeysetHandle,
    val wrapped: RecoveryKeyManager.WrappedKey,
    val recoveryPhrase: String
)

/**
 * High-level E2EE orchestration injected into repositories and ViewModels.
 * All KDF operations run on [Dispatchers.Default] to keep the main thread free.
 */
@Singleton
class CoupleEncryptionManager @Inject constructor(
    private val keyStore: CoupleKeyStore,
    private val keyManager: RecoveryKeyManager
) {
    /**
     * Generates a new couple keyset + recovery phrase but does NOT store them yet —
     * storage is deferred until the server confirms the invite code via [storeInviteSetup].
     * This avoids storing the keyset under a client-generated code that may differ from
     * the code the server ultimately uses (which would break reconciliation).
     */
    suspend fun setupForNewCouple(): SetupResult = withContext(Dispatchers.Default) {
        val phrase = keyManager.generateRecoveryPhrase()
        val handle = keyManager.newCoupleKeyset()
        val wrapped = keyManager.wrap(handle, phrase)
        SetupResult(handle, wrapped, phrase)
    }

    /** Persists a [SetupResult] under the server-confirmed [inviteCode]. */
    fun storeInviteSetup(inviteCode: String, result: SetupResult) {
        keyStore.storeInviteKeyset(inviteCode, result.handle)
        keyStore.storeInvitePhrase(inviteCode, result.recoveryPhrase)
    }

    /**
     * Called by the acceptor during pairing, and on any device during recovery.
     * Derives the wrap key from [phrase], unwraps the keyset, and stores it
     * locally under [coupleId]. Also persists [phrase] so it can be shown in settings.
     * Throws [com.google.crypto.tink.subtle.Validators]-wrapped exception if phrase is wrong.
     */
    suspend fun unwrapAndStore(
        coupleId: String,
        wrapped: RecoveryKeyManager.WrappedKey,
        phrase: String
    ): Result<Unit> = withContext(Dispatchers.Default) {
        runCatching {
            val handle = keyManager.unwrap(wrapped, phrase)
            keyStore.storeKeyset(coupleId, handle)
            keyStore.storeRecoveryPhrase(coupleId, phrase)
        }
    }

    fun recoveryPhrase(coupleId: String): String? = keyStore.loadRecoveryPhrase(coupleId)

    /**
     * Called on app launch / Home load after the couple doc is resolved. Every couple is
     * strict-E2EE, so the only question is whether this device holds the couple key:
     * present -> UNLOCKED; recoverable from the inviter's invite slot -> RECONCILED_FROM_INVITE;
     * otherwise -> NEEDS_RECOVERY (prompt for the recovery phrase).
     */
    fun checkStatus(couple: Couple): EncryptionStatus {
        if (keyStore.hasKeyset(couple.id)) return EncryptionStatus.UNLOCKED
        if (keyStore.reconcileInviteKeyset(couple.inviteCode, couple.id)) {
            return EncryptionStatus.RECONCILED_FROM_INVITE
        }
        return EncryptionStatus.NEEDS_RECOVERY
    }

    fun aeadFor(coupleId: String): Aead? = keyStore.aeadFor(coupleId)

    /** Answer-bearing writes must never fall back to plaintext. */
    fun requireAead(coupleId: String): Aead =
        keyStore.aeadFor(coupleId) ?: throw MissingCoupleKeyException(coupleId)

    fun isUnlocked(coupleId: String): Boolean = keyStore.hasKeyset(coupleId)

    /**
     * Re-wraps the locally-held keyset with a new phrase and returns the new WrappedKey
     * so the caller can persist it to Firestore. The old phrase is NOT required.
     * Returns failure if no local keyset exists for [coupleId].
     */
    suspend fun rewrapWithNewPhrase(
        coupleId: String,
        newPhrase: String
    ): Result<RecoveryKeyManager.WrappedKey> = withContext(Dispatchers.Default) {
        runCatching {
            val handle = keyStore.loadKeyset(coupleId)
                ?: error("No local keyset for $coupleId -- cannot change phrase without recovery first")
            keyManager.wrap(handle, newPhrase)
        }
    }

    fun deleteKeyset(coupleId: String) = keyStore.deleteKeyset(coupleId)
}
