package app.closer.crypto

import app.closer.crypto.EncryptionVersion.STRICT
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. */
    RECONCILED_FROM_INVITE,
    /** encryptionVersion == 1 but no local keyset -- prompt for recovery phrase. */
    NEEDS_RECOVERY,
    /** encryptionVersion == 0 -- this couple must create a key before writing more answers. */
    NEEDS_ENCRYPTION_UPGRADE,
    /** encryptionVersion == 1 with a local key -- this device must rewrite its legacy answers. */
    NEEDS_CONTENT_MIGRATION
}

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.
     * Handles inviter reconciliation (flow B') transparently.
     */
    fun checkStatus(couple: Couple): EncryptionStatus {
        // v2 couples were created by Android with a strict couple key.
        // v1 couples are mid-migration; v0 couples are plaintext (iOS MVP).
        when (couple.encryptionVersion) {
            EncryptionVersion.PLAINTEXT -> return EncryptionStatus.NEEDS_ENCRYPTION_UPGRADE
            EncryptionVersion.MIGRATING -> { /* fall through to keyset checks below */ }
            EncryptionVersion.STRICT -> { /* fall through to keyset checks below */ }
        }
        if (keyStore.hasKeyset(couple.id)) {
            return if (couple.encryptionVersion >= STRICT) {
                EncryptionStatus.UNLOCKED
            } else {
                EncryptionStatus.NEEDS_CONTENT_MIGRATION
            }
        }
        if (keyStore.reconcileInviteKeyset(couple.inviteCode, couple.id)) {
            return if (couple.encryptionVersion >= STRICT) {
                EncryptionStatus.RECONCILED_FROM_INVITE
            } else {
                EncryptionStatus.NEEDS_CONTENT_MIGRATION
            }
        }
        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)

    suspend fun setupLegacyCouple(coupleId: String): SetupResult = withContext(Dispatchers.Default) {
        val phrase = keyManager.generateRecoveryPhrase()
        val handle = keyManager.newCoupleKeyset()
        val wrapped = keyManager.wrap(handle, phrase)
        keyStore.storeKeyset(coupleId, handle)
        keyStore.storePendingRecoveryPhrase(coupleId, phrase)
        SetupResult(handle, wrapped, phrase)
    }

    fun pendingRecoveryPhrase(coupleId: String): String? =
        keyStore.pendingRecoveryPhrase(coupleId)

    fun acknowledgeRecoveryPhrase(coupleId: String) =
        keyStore.clearPendingRecoveryPhrase(coupleId)

    companion object {
        /** Kept for backwards compatibility; prefer [EncryptionVersion.STRICT]. */
        const val STRICT_ENCRYPTION_VERSION = EncryptionVersion.STRICT
    }
}
