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 org.json.JSONObject
import javax.inject.Inject
import javax.inject.Singleton

/** One answer in a How Well session — either a chosen option or a scale value. */
data class HowWellRawAnswer(val optionId: String? = null, val scale: Int? = null)

/** Both partners' answers for one How Well session, keyed by userId. */
data class HowWellAnswers(
    val byUser: Map<String, List<HowWellRawAnswer>> = emptyMap()
)

/**
 * Stores each partner's answers for the async How Well reveal at
 * `couples/{coupleId}/how_well/{sessionId}`. The subject (session starter) answers
 * honestly; the guesser predicts. The session (shared question set + the
 * one-active-game lock) is owned by [app.closer.domain.usecase.GameSessionManager];
 * this carries the answers so each partner plays on their own device and the score
 * reveals once both have submitted.
 */
@Singleton
class FirestoreHowWellDataSource @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.HOW_WELL)
            .document(sessionId)

    /** Persist this user's full answer list (aligned to the session's question order). */
    suspend fun submitAnswers(
        coupleId: String,
        sessionId: String,
        userId: String,
        answers: List<HowWellRawAnswer>
    ) {
        val aead = encryptionManager.requireAead(coupleId)
        val json = JSONArray(answers.map {
            JSONObject().apply {
                put("optionId", it.optionId ?: JSONObject.NULL)
                put("scale", it.scale ?: JSONObject.NULL)
            }
        }.toString())
        val value = fieldEncryptor.encrypt(json.toString(), aead, coupleId)
        doc(coupleId, sessionId)
            .set(mapOf("answers" to mapOf(userId to value)), SetOptions.merge())
            .await()
    }

    /** One-shot read — used to detect whether this user has already answered. */
    suspend fun getAnswers(coupleId: String, sessionId: String): HowWellAnswers? =
        runCatching {
            val aead = encryptionManager.aeadFor(coupleId)
            HowWellAnswers(parseAnswers(doc(coupleId, sessionId).get().await().get("answers"), aead, coupleId))
        }.getOrNull()

    /** Live view of both partners' answers; emits whenever either side submits. */
    fun observeAnswers(coupleId: String, sessionId: String): Flow<HowWellAnswers> =
        callbackFlow {
            val reg = doc(coupleId, sessionId).addSnapshotListener { snap, err ->
                if (err != null || snap == null) return@addSnapshotListener
                val aead = encryptionManager.aeadFor(coupleId)
                trySend(HowWellAnswers(parseAnswers(snap.get("answers"), aead, coupleId)))
            }
            awaitClose { reg.remove() }
        }

    private fun parseAnswers(
        raw: Any?,
        aead: com.google.crypto.tink.Aead?,
        coupleId: String
    ): Map<String, List<HowWellRawAnswer>> {
        @Suppress("UNCHECKED_CAST")
        val map = raw as? Map<String, *> ?: return emptyMap()
        return map.mapNotNull { (uid, value) ->
            val list = when (value) {
                is String -> listOf(value)
                is List<*> -> value
                else -> return@mapNotNull null
            }
            // Encrypted: single-element list with JSON blob
            if (list.size == 1 && list[0] is String && fieldEncryptor.isEncrypted(list[0] as String)) {
                val json = fieldEncryptor.decrypt(list[0] as String, aead, coupleId) ?: return@mapNotNull null
                val answers = runCatching {
                    val arr = JSONArray(json)
                    (0 until arr.length()).map { i ->
                        val obj = arr.getJSONObject(i)
                        HowWellRawAnswer(
                            optionId = if (obj.isNull("optionId")) null else obj.getString("optionId"),
                            scale = if (obj.isNull("scale")) null else obj.getInt("scale")
                        )
                    }
                }.getOrNull() ?: return@mapNotNull null
                uid to answers
            } else {
                uid to list.mapNotNull { item ->
                    (item as? Map<*, *>)?.let {
                        HowWellRawAnswer(
                            optionId = it["optionId"] as? String,
                            scale = (it["scale"] as? Number)?.toInt()
                        )
                    }
                }
            }
        }.toMap()
    }
}
