package app.closer.data.remote

import app.closer.crypto.AnswerCommitment
import app.closer.crypto.CoupleEncryptionManager
import app.closer.crypto.FieldEncryptor
import app.closer.crypto.PendingAnswerKeyStore
import app.closer.crypto.SealedAnswerEncryptor
import app.closer.crypto.UserKeyManager
import app.closer.domain.model.QuestionAnswer
import app.closer.domain.model.QuestionMessage
import app.closer.domain.model.QuestionReaction
import app.closer.domain.model.QuestionThread
import app.closer.domain.model.QuestionThreadStatus
import org.json.JSONArray
import com.google.firebase.firestore.DocumentSnapshot
import com.google.firebase.firestore.FieldValue
import com.google.firebase.firestore.FirebaseFirestore
import com.google.firebase.firestore.Query
import kotlinx.coroutines.channels.awaitClose
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.callbackFlow
import kotlinx.coroutines.suspendCancellableCoroutine
import javax.inject.Inject
import javax.inject.Singleton
import kotlin.coroutines.resume
import kotlin.coroutines.resumeWithException

@Singleton
class FirestoreQuestionThreadDataSource @Inject constructor(
    private val db: FirebaseFirestore,
    private val encryptionManager: CoupleEncryptionManager,
    private val fieldEncryptor: FieldEncryptor,
    private val userKeyManager: UserKeyManager,
    private val deviceKeyDataSource: FirestoreDeviceKeyDataSource,
    private val sealedAnswerEncryptor: SealedAnswerEncryptor,
    private val pendingAnswerKeyStore: PendingAnswerKeyStore,
    private val answerCommitment: AnswerCommitment,
    private val storageDataSource: FirebaseStorageDataSource
) {

    private fun threadsRef(coupleId: String) =
        db.collection(FirestoreCollections.COUPLES).document(coupleId)
            .collection(FirestoreCollections.Couples.QUESTION_THREADS)

    // ─── Thread ─────────────────────────────────────────────────────────────────

    suspend fun findThreadByQuestionId(coupleId: String, questionId: String): String? {
        val snap = threadsRef(coupleId)
            .whereEqualTo("questionId", questionId)
            .limit(1)
            .getAwait()
        return snap.documents.firstOrNull()?.id
    }

    suspend fun createThread(coupleId: String, questionId: String, categoryId: String, userId: String): String {
        val now = FieldValue.serverTimestamp()
        val doc = threadsRef(coupleId).document()
        doc.set(
            mapOf(
                "questionId" to questionId,
                "categoryId" to categoryId,
                "status" to QuestionThreadStatus.NOT_STARTED.toFirestoreValue(),
                "currentIndex" to 0,
                "createdByUserId" to userId,
                "createdAt" to now,
                "updatedAt" to now
            )
        ).voidAwait()
        return doc.id
    }

    fun observeThread(coupleId: String, threadId: String): Flow<QuestionThread> = callbackFlow {
        val listener = threadsRef(coupleId).document(threadId)
            .addSnapshotListener { snap, err ->
                if (err != null || snap == null) return@addSnapshotListener
                trySend(snap.toQuestionThread(coupleId))
            }
        awaitClose { listener.remove() }
    }

    suspend fun updateThreadStatus(coupleId: String, threadId: String, status: QuestionThreadStatus) {
        threadsRef(coupleId).document(threadId)
            .update("status", status.toFirestoreValue())
            .voidAwait()
    }

    // ─── Answers ─────────────────────────────────────────────────────────────────

    suspend fun submitAnswer(coupleId: String, threadId: String, userId: String, answer: QuestionAnswer) {
        submitAnswerSealed(coupleId, threadId, userId, answer)
    }

    // schemaVersion 3: per-answer one-time key — partner-proof before reveal.
    // threadId is used as the AAD "questionId" so thread keys are distinct from
    // daily-question keys even when the same question appears in both contexts.
    private suspend fun submitAnswerSealed(coupleId: String, threadId: String, userId: String, answer: QuestionAnswer) {
        ensureUserPublicKeyPublished(userId)
        val oneTimeKey = sealedAnswerEncryptor.generateOneTimeKey()
        val payload = SealedAnswerEncryptor.AnswerPayload(
            writtenText = answer.writtenText,
            selectedOptionIds = answer.selectedOptionIds,
            scaleValue = answer.scaleValue
        )
        val encryptedPayload = sealedAnswerEncryptor.seal(payload, oneTimeKey, coupleId, threadId, userId)
        val commitment = answerCommitment.compute(
            coupleId, threadId, userId,
            answer.writtenText, answer.selectedOptionIds, answer.scaleValue
        )
        pendingAnswerKeyStore.store("thread_$threadId", oneTimeKey)

        threadsRef(coupleId).document(threadId)
            .collection(FirestoreCollections.QuestionThreads.ANSWERS)
            .document(userId)
            .set(
                mapOf(
                    "userId" to userId,
                    "questionId" to answer.questionId,
                    "answerType" to answer.answerType,
                    "encryptedPayload" to encryptedPayload,
                    "commitmentHash" to commitment,
                    "schemaVersion" to 3,
                    "answerKeyReleased" to false,
                    "createdAt" to FieldValue.serverTimestamp(),
                    "updatedAt" to FieldValue.serverTimestamp()
                )
            ).voidAwait()
    }

    // Call after releasing the one-time key so the answer doc reflects the released state.
    // Required for correct phase detection on cold restart of the reveal screen.
    suspend fun markAnswerKeyReleased(coupleId: String, threadId: String, userId: String) {
        threadsRef(coupleId).document(threadId)
            .collection(FirestoreCollections.QuestionThreads.ANSWERS)
            .document(userId)
            .update(mapOf("answerKeyReleased" to true, "updatedAt" to FieldValue.serverTimestamp()))
            .voidAwait()
    }

    suspend fun getAnswerCount(coupleId: String, threadId: String): Int {
        val snap = threadsRef(coupleId)
            .document(threadId)
            .collection(FirestoreCollections.QuestionThreads.ANSWERS)
            .getAwait()
        return snap.size()
    }

    fun observeAnswers(coupleId: String, threadId: String): Flow<List<QuestionAnswer>> = callbackFlow {
        val listener = threadsRef(coupleId)
            .document(threadId)
            .collection(FirestoreCollections.QuestionThreads.ANSWERS)
            .addSnapshotListener { snap, err ->
                if (err != null || snap == null) return@addSnapshotListener
                trySend(snap.documents.mapNotNull { it.toQuestionAnswer() })
            }
        awaitClose { listener.remove() }
    }

    // ─── Messages ────────────────────────────────────────────────────────────────

    suspend fun sendMessage(coupleId: String, threadId: String, message: QuestionMessage) {
        val aead = encryptionManager.requireAead(coupleId)
        threadsRef(coupleId)
            .document(threadId)
            .collection(FirestoreCollections.QuestionThreads.MESSAGES)
            .add(
                mapOf(
                    "authorUserId" to message.userId,
                    "type" to "text",
                    "text" to fieldEncryptor.encrypt(message.text, aead, coupleId),
                    "createdAt" to FieldValue.serverTimestamp()
                )
            ).refAwait()
    }

    /**
     * Sends an image message: the bytes are encrypted with the couple key on-device, the ciphertext
     * is uploaded to Storage, and only the (encrypted) media's URL is stored in Firestore.
     */
    suspend fun sendImageMessage(coupleId: String, threadId: String, userId: String, imageBytes: ByteArray) {
        val aead = encryptionManager.requireAead(coupleId)
        val encrypted = aead.encrypt(imageBytes, coupleId.toByteArray(Charsets.UTF_8))
        val url = storageDataSource.uploadEncryptedMedia(userId, encrypted)
        threadsRef(coupleId)
            .document(threadId)
            .collection(FirestoreCollections.QuestionThreads.MESSAGES)
            .add(
                mapOf(
                    "authorUserId" to userId,
                    "type" to "image",
                    "mediaUrl" to url,
                    "createdAt" to FieldValue.serverTimestamp()
                )
            ).refAwait()
    }

    /** Downloads + decrypts an image message's bytes for display (couple key, on-device). */
    suspend fun loadDecryptedMedia(coupleId: String, mediaUrl: String): ByteArray? {
        val aead = encryptionManager.aeadFor(coupleId) ?: return null
        val cipher = runCatching { storageDataSource.downloadBytes(mediaUrl) }.getOrNull() ?: return null
        return runCatching { aead.decrypt(cipher, coupleId.toByteArray(Charsets.UTF_8)) }.getOrNull()
    }

    fun observeMessages(coupleId: String, threadId: String): Flow<List<QuestionMessage>> = callbackFlow {
        val listener = threadsRef(coupleId)
            .document(threadId)
            .collection(FirestoreCollections.QuestionThreads.MESSAGES)
            .orderBy("createdAt", Query.Direction.ASCENDING)
            .addSnapshotListener { snap, err ->
                if (err != null || snap == null) return@addSnapshotListener
                val aead = encryptionManager.aeadFor(coupleId)
                trySend(snap.documents.mapNotNull { it.toQuestionMessage(aead, coupleId) })
            }
        awaitClose { listener.remove() }
    }

    // ─── Reactions ───────────────────────────────────────────────────────────────

    suspend fun addReaction(coupleId: String, threadId: String, reaction: QuestionReaction) {
        val docId = "${reaction.userId}_${reaction.targetUserId}"
        threadsRef(coupleId)
            .document(threadId)
            .collection(FirestoreCollections.QuestionThreads.REACTIONS)
            .document(docId)
            .set(
                mapOf(
                    "userId" to reaction.userId,
                    "targetUserId" to reaction.targetUserId,
                    "emoji" to reaction.emoji,
                    "createdAt" to FieldValue.serverTimestamp()
                )
            ).voidAwait()
    }

    fun observeReactions(coupleId: String, threadId: String): Flow<List<QuestionReaction>> = callbackFlow {
        val listener = threadsRef(coupleId)
            .document(threadId)
            .collection(FirestoreCollections.QuestionThreads.REACTIONS)
            .addSnapshotListener { snap, err ->
                if (err != null || snap == null) return@addSnapshotListener
                trySend(snap.documents.mapNotNull { it.toQuestionReaction() })
            }
        awaitClose { listener.remove() }
    }

    // ─── Coroutine helpers ───────────────────────────────────────────────────────

    private suspend fun com.google.firebase.firestore.CollectionReference.getAwait() =
        get().queryAwait()

    private suspend fun com.google.firebase.firestore.Query.getAwait() =
        get().queryAwait()

    private suspend fun com.google.android.gms.tasks.Task<com.google.firebase.firestore.QuerySnapshot>.queryAwait() =
        suspendCancellableCoroutine { cont ->
            addOnSuccessListener { cont.resume(it) }
            addOnFailureListener { cont.resumeWithException(it) }
        }

    private suspend fun com.google.android.gms.tasks.Task<Void>.voidAwait() =
        suspendCancellableCoroutine<Unit> { cont ->
            addOnSuccessListener { cont.resume(Unit) }
            addOnFailureListener { cont.resumeWithException(it) }
        }

    private suspend fun com.google.android.gms.tasks.Task<com.google.firebase.firestore.DocumentReference>.refAwait() =
        suspendCancellableCoroutine<Unit> { cont ->
            addOnSuccessListener { cont.resume(Unit) }
            addOnFailureListener { cont.resumeWithException(it) }
        }

    // ─── Document mappers ────────────────────────────────────────────────────────

    private fun DocumentSnapshot.toQuestionThread(coupleId: String) = QuestionThread(
        id = id,
        coupleId = coupleId,
        questionId = getString("questionId") ?: "",
        categoryId = getString("categoryId") ?: "",
        status = QuestionThreadStatus.fromFirestoreValue(getString("status") ?: ""),
        currentIndex = getLong("currentIndex")?.toInt() ?: 0,
        createdAt = getTimestamp("createdAt")?.toDate()?.time ?: 0L,
        updatedAt = getTimestamp("updatedAt")?.toDate()?.time ?: 0L
    )

    private fun DocumentSnapshot.toQuestionAnswer(): QuestionAnswer? {
        val userId = getString("userId") ?: return null
        // All thread answers are sealed (schemaVersion 3): content lives in encryptedPayload
        // and is decrypted by the reveal flow, never stored in plaintext.
        return QuestionAnswer(
            userId = userId,
            questionId = getString("questionId") ?: "",
            answerType = getString("answerType") ?: "written",
            schemaVersion = 3,
            isSealed = getBoolean("answerKeyReleased") != true,
            encryptedPayload = getString("encryptedPayload"),
            createdAt = getTimestamp("createdAt")?.toDate()?.time ?: 0L,
            updatedAt = getTimestamp("updatedAt")?.toDate()?.time ?: 0L
        )
    }

    private fun DocumentSnapshot.toQuestionMessage(
        aead: com.google.crypto.tink.Aead?,
        coupleId: String
    ): QuestionMessage? {
        val userId = getString("authorUserId") ?: return null
        val type = getString("type") ?: "text"
        return QuestionMessage(
            id = id,
            userId = userId,
            type = type,
            mediaUrl = getString("mediaUrl") ?: "",
            text = if (type == "image") "" else (fieldEncryptor.decryptForDisplay(getString("text"), aead, coupleId) ?: ""),
            createdAt = getTimestamp("createdAt")?.toDate()?.time ?: 0L
        )
    }

    private fun DocumentSnapshot.toQuestionReaction(): QuestionReaction? {
        val userId = getString("userId") ?: return null
        return QuestionReaction(
            id = id,
            userId = userId,
            targetUserId = getString("targetUserId") ?: "",
            emoji = getString("emoji") ?: "",
            createdAt = getTimestamp("createdAt")?.toDate()?.time ?: 0L
        )
    }

    private suspend fun ensureUserPublicKeyPublished(userId: String) {
        if (deviceKeyDataSource.getPublicKey(userId) != null) return
        val privateKey = userKeyManager.getOrCreatePrivateKey()
        deviceKeyDataSource.publishPublicKey(userId, userKeyManager.publicKeyB64(privateKey))
    }
}
