package app.closer.crypto

import app.closer.data.remote.FirestoreDeviceKeyDataSource
import app.closer.data.remote.FirestoreReleaseKeyDataSource
import javax.inject.Inject
import javax.inject.Singleton

/**
 * Orchestrates the two-sided reveal key exchange for sealed answers.
 *
 * Reveal flow (called once both answer docs exist):
 *  1. [releaseOwnKey]: encrypt our pending one-time key to the partner's public key,
 *     write the keybox to Firestore. After this, the partner can decrypt our answer.
 *  2. [decryptPartnerAnswer]: read the keybox the partner wrote for us, unwrap our copy
 *     of their one-time key, and decrypt their sealed answer payload.
 *
 * A user cannot see the partner's answer until their own key is released (step 1 must
 * succeed before step 2 is attempted). This prevents one-sided reveal abuse.
 */
@Singleton
class SealedRevealManager @Inject constructor(
    private val userKeyManager: UserKeyManager,
    private val pendingAnswerKeyStore: PendingAnswerKeyStore,
    private val releaseKeyEncryptor: ReleaseKeyEncryptor,
    private val sealedAnswerEncryptor: SealedAnswerEncryptor,
    private val deviceKeyDataSource: FirestoreDeviceKeyDataSource,
    private val releaseKeyDataSource: FirestoreReleaseKeyDataSource
) {

    /**
     * Releases the user's own one-time answer key to their partner.
     *
     * @param coupleId The couple identifier.
     * @param date The daily-question date string (YYYY-MM-DD).
     * @param questionId The question ID (used as AAD).
     * @param userId The current user's UID.
     * @param partnerId The partner's UID.
     * @return true if the key was released successfully, false if the pending key is missing.
     */
    suspend fun releaseOwnKey(
        coupleId: String,
        date: String,
        questionId: String,
        userId: String,
        partnerId: String
    ): Boolean {
        val storeKey = "daily_${date}_${questionId}"
        val oneTimeKey = pendingAnswerKeyStore.load(storeKey) ?: return false
        val partnerPublicKey = deviceKeyDataSource.getPublicKey(partnerId) ?: return false

        val keybox = releaseKeyEncryptor.wrapForRecipient(
            oneTimeKey = oneTimeKey,
            recipientPublicKeyB64 = partnerPublicKey,
            coupleId = coupleId,
            questionId = questionId,
            senderUserId = userId,
            recipientUserId = partnerId
        )

        releaseKeyDataSource.writeReleaseKey(
            coupleId = coupleId,
            date = date,
            senderUserId = userId,
            recipientUserId = partnerId,
            encryptedAnswerKey = keybox
        )

        pendingAnswerKeyStore.remove(storeKey)
        return true
    }

    /**
     * Decrypts the partner's sealed answer after they have released their key to us.
     *
     * @return The decrypted [SealedAnswerEncryptor.AnswerPayload], or null if the
     *   partner has not yet released their key.
     */
    suspend fun decryptPartnerAnswer(
        coupleId: String,
        date: String,
        questionId: String,
        partnerId: String,
        userId: String,
        encryptedPayload: String
    ): SealedAnswerEncryptor.AnswerPayload? {
        val keybox = releaseKeyDataSource.readReleaseKey(
            coupleId = coupleId,
            date = date,
            senderUserId = partnerId,
            recipientUserId = userId
        ) ?: return null

        // loadPrivateKey — NOT getOrCreatePrivateKey. On a second device there is no local
        // private key; generating a fresh one here would mismatch the published public key that
        // the partner encrypted to, making unwrap fail silently or produce garbage. Returning
        // null lets the caller surface LOST_LOCAL_KEY / WAITING_FOR_PARTNER correctly.
        val myPrivateKey = userKeyManager.loadPrivateKey() ?: return null
        val oneTimeKey = releaseKeyEncryptor.unwrapFromSender(
            keyboxB64 = keybox,
            recipientPrivateKey = myPrivateKey,
            coupleId = coupleId,
            questionId = questionId,
            senderUserId = partnerId,
            recipientUserId = userId
        )

        return sealedAnswerEncryptor.open(
            encryptedPayload = encryptedPayload,
            keyHandle = oneTimeKey,
            coupleId = coupleId,
            questionId = questionId,
            userId = partnerId
        )
    }

    // ── Thread reveal ─────────────────────────────────────────────────────────────
    // Thread answers use the threadId (not a date) as the key-store identifier and
    // as the AAD "questionId", so daily-question and thread keys for the same
    // question ID are always cryptographically distinct.

    suspend fun releaseOwnKeyForThread(
        coupleId: String,
        threadId: String,
        userId: String,
        partnerId: String
    ): Boolean {
        val storeKey = "thread_$threadId"
        val oneTimeKey = pendingAnswerKeyStore.load(storeKey) ?: return false
        val partnerPublicKey = deviceKeyDataSource.getPublicKey(partnerId) ?: return false

        val keybox = releaseKeyEncryptor.wrapForRecipient(
            oneTimeKey = oneTimeKey,
            recipientPublicKeyB64 = partnerPublicKey,
            coupleId = coupleId,
            questionId = threadId,
            senderUserId = userId,
            recipientUserId = partnerId
        )

        releaseKeyDataSource.writeReleaseKeyForThread(
            coupleId = coupleId,
            threadId = threadId,
            senderUserId = userId,
            recipientUserId = partnerId,
            encryptedAnswerKey = keybox
        )

        // Unlike daily answers (whose plaintext is mirrored to local storage), a thread answer's
        // only plaintext source on this device is this one-time key. Keep it so the reveal screen
        // can re-decrypt our own answer on cold restart.
        return true
    }

    /** Decrypts the user's OWN sealed thread answer using their locally-held one-time key. */
    suspend fun decryptOwnThreadAnswer(
        coupleId: String,
        threadId: String,
        userId: String,
        encryptedPayload: String
    ): SealedAnswerEncryptor.AnswerPayload? {
        val oneTimeKey = pendingAnswerKeyStore.load("thread_$threadId") ?: return null
        return sealedAnswerEncryptor.open(
            encryptedPayload = encryptedPayload,
            keyHandle = oneTimeKey,
            coupleId = coupleId,
            questionId = threadId,
            userId = userId
        )
    }

    suspend fun decryptPartnerThreadAnswer(
        coupleId: String,
        threadId: String,
        partnerId: String,
        userId: String,
        encryptedPayload: String
    ): SealedAnswerEncryptor.AnswerPayload? {
        val keybox = releaseKeyDataSource.readReleaseKeyForThread(
            coupleId = coupleId,
            threadId = threadId,
            senderUserId = partnerId,
            recipientUserId = userId
        ) ?: return null

        val myPrivateKey = userKeyManager.loadPrivateKey() ?: return null
        val oneTimeKey = releaseKeyEncryptor.unwrapFromSender(
            keyboxB64 = keybox,
            recipientPrivateKey = myPrivateKey,
            coupleId = coupleId,
            questionId = threadId,
            senderUserId = partnerId,
            recipientUserId = userId
        )

        return sealedAnswerEncryptor.open(
            encryptedPayload = encryptedPayload,
            keyHandle = oneTimeKey,
            coupleId = coupleId,
            questionId = threadId,
            userId = partnerId
        )
    }

    /**
     * Ensures this user's public key is published to Firestore.
     * Safe to call on every launch — no-ops if already published.
     */
    suspend fun ensurePublicKeyPublished(userId: String) {
        val existing = deviceKeyDataSource.getPublicKey(userId)
        if (existing != null) return
        val privateKey = userKeyManager.getOrCreatePrivateKey()
        val publicKeyB64 = userKeyManager.publicKeyB64(privateKey)
        deviceKeyDataSource.publishPublicKey(userId, publicKeyB64)
    }
}
