package app.closer.data.remote

import app.closer.crypto.CoupleEncryptionManager
import app.closer.crypto.FieldEncryptor
import com.google.firebase.firestore.FirebaseFirestore
import com.google.firebase.firestore.SetOptions
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.callbackFlow
import kotlinx.coroutines.tasks.await
import org.json.JSONArray
import javax.inject.Inject
import javax.inject.Singleton

/**
 * Both partners' picks for one This-or-That session, keyed by userId. Each list
 * holds the chosen optionIds aligned to the session's `questionIds` order, so
 * `byUser[a][i]` and `byUser[b][i]` are answers to the same prompt.
 */
data class ThisOrThatAnswers(
    val byUser: Map<String, List<String>> = emptyMap()
)

/**
 * Stores per-user answers for the async This-or-That reveal at
 * `couples/{coupleId}/this_or_that/{sessionId}`. The session itself (the shared
 * question set + the one-active-game lock) is owned by [app.closer.domain.usecase.GameSessionManager];
 * this only carries the answers so each partner can play on their own device and
 * the result reveals once both have submitted.
 */
@Singleton
class FirestoreThisOrThatDataSource @Inject constructor(
    private val db: FirebaseFirestore,
    private val encryptionManager: CoupleEncryptionManager,
    private val fieldEncryptor: FieldEncryptor
) {
    private fun doc(coupleId: String, sessionId: String) =
        db.collection(FirestoreCollections.COUPLES)
            .document(coupleId)
            .collection(FirestoreCollections.Couples.THIS_OR_THAT)
            .document(sessionId)

    /** Persist this user's full set of picks (optionIds in `questionIds` order). */
    suspend fun submitAnswers(
        coupleId: String,
        sessionId: String,
        userId: String,
        optionIds: List<String>
    ) {
        val aead = encryptionManager.requireAead(coupleId)
        val encrypted = fieldEncryptor.encrypt(JSONArray(optionIds).toString(), aead, coupleId)
        doc(coupleId, sessionId)
            .set(mapOf("answers" to mapOf(userId to encrypted)), SetOptions.merge())
            .await()
    }

    /** One-shot load of both partners' picks (for history replay). */
    suspend fun getAnswers(coupleId: String, sessionId: String): ThisOrThatAnswers? =
        runCatching {
            val snap = doc(coupleId, sessionId).get().await()
            if (!snap.exists()) return@runCatching null
            @Suppress("UNCHECKED_CAST")
            val raw = snap.get("answers") as? Map<String, *>
            val byUser = parseAnswers(raw, coupleId)
            ThisOrThatAnswers(byUser)
        }.getOrNull()

    /** Live view of both partners' picks; emits whenever either side submits. */
    fun observeAnswers(coupleId: String, sessionId: String): Flow<ThisOrThatAnswers> =
        callbackFlow {
            val reg = doc(coupleId, sessionId).addSnapshotListener { snap, err ->
                if (err != null || snap == null) return@addSnapshotListener
                @Suppress("UNCHECKED_CAST")
                val raw = snap.get("answers") as? Map<String, *>
                val byUser = parseAnswers(raw, coupleId)
                trySend(ThisOrThatAnswers(byUser))
            }
            awaitClose { reg.remove() }
        }

    private fun parseAnswers(raw: Map<String, *>?, coupleId: String): Map<String, List<String>> {
        val aead = encryptionManager.aeadFor(coupleId)
        return raw.orEmpty().mapNotNull { (uid, value) ->
            when (value) {
                is String -> {
                    val json = fieldEncryptor.decrypt(value, aead, coupleId) ?: return@mapNotNull null
                    val answers = runCatching {
                        val array = JSONArray(json)
                        (0 until array.length()).map { array.getString(it) }
                    }.getOrNull() ?: return@mapNotNull null
                    uid to answers
                }
                // Version-0 compatibility exists only until this user completes migration.
                is List<*> -> uid to value.filterIsInstance<String>()
                else -> null
            }
        }.toMap()
    }
}
