package app.closer.ui.answers

import androidx.lifecycle.SavedStateHandle
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import app.closer.analytics.RetentionAnalytics
import app.closer.analytics.RetentionEvent
import app.closer.core.navigation.AppRoute
import app.closer.core.crash.CrashReporter
import app.closer.crypto.PendingAnswerKeyStore
import app.closer.crypto.SealedRevealManager
import app.closer.data.remote.FirestoreAnswerDataSource
import app.closer.domain.model.LocalAnswer
import app.closer.domain.model.Question
import app.closer.domain.repository.AuthRepository
import app.closer.domain.repository.CoupleRepository
import app.closer.domain.repository.LocalAnswerRepository
import app.closer.domain.repository.QuestionRepository
import dagger.hilt.android.lifecycle.HiltViewModel
import javax.inject.Inject
import kotlinx.coroutines.Job
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch

enum class FollowUpOption(val label: String, val route: String? = null) {
    DEEPER_FOLLOW_UP("Want to ask one deeper follow-up?", AppRoute.QUESTION_COMPOSER),
    DATE_IDEA("Want to turn this into a date idea?", AppRoute.DATE_BUILDER),
    SAVE_MEMORY("Want to save this as a memory?", AppRoute.MEMORY_LANE),
    ANOTHER_QUESTION("Want to try another question from this category?")
}

/**
 * Tracks where we are in the sealed-answer reveal exchange.
 *
 * State machine:
 *   NONE            — answer is enc:v1: (schemaVersion 2); use existing non-sealed reveal path
 *   ANSWER_SEALED   — own answer submitted; partner has not answered yet
 *   BOTH_ANSWERED   — both answers submitted; "Reveal is ready" but keys not released yet
 *   RELEASING_KEY   — writing our one-time key to Firestore (in-flight)
 *   WAITING_FOR_PARTNER — our key released; partner has not released theirs yet
 *   REVEALED        — both keys released, partner answer decrypted and visible
 *   LOST_LOCAL_KEY  — the pending answer key is missing from this device (unrecoverable)
 */
enum class SealedRevealPhase {
    NONE,
    ANSWER_SEALED,
    BOTH_ANSWERED,
    RELEASING_KEY,
    WAITING_FOR_PARTNER,
    REVEALED,
    LOST_LOCAL_KEY
}

data class AnswerRevealUiState(
    val isLoading: Boolean = true,
    val error: String? = null,
    val question: Question? = null,
    val answer: LocalAnswer? = null,
    val partnerAnswer: LocalAnswer? = null,
    val coupleId: String? = null,
    val partnerId: String? = null,
    val followUpOptions: List<FollowUpOption> = emptyList(),
    val snackbarMessage: String? = null,
    val sealedRevealPhase: SealedRevealPhase = SealedRevealPhase.NONE,
    val loreSaved: Boolean = false
)

@HiltViewModel
class AnswerRevealViewModel @Inject constructor(
    private val questionRepository: QuestionRepository,
    private val localAnswerRepository: LocalAnswerRepository,
    private val firestoreAnswerDataSource: FirestoreAnswerDataSource,
    private val authRepository: AuthRepository,
    private val coupleRepository: CoupleRepository,
    private val crashReporter: CrashReporter,
    private val sealedRevealManager: SealedRevealManager,
    private val pendingAnswerKeyStore: PendingAnswerKeyStore,
    private val retentionAnalytics: RetentionAnalytics,
    savedStateHandle: SavedStateHandle
) : ViewModel() {

    private val questionId: String = savedStateHandle["questionId"] ?: ""

    private val _uiState = MutableStateFlow(AnswerRevealUiState())
    val uiState: StateFlow<AnswerRevealUiState> = _uiState.asStateFlow()

    private var partnerObserveJob: Job? = null

    init {
        load()
        observeAnswer()
    }

    private fun load() {
        viewModelScope.launch {
            _uiState.value = AnswerRevealUiState(isLoading = true)
            try {
                val (coupleId, partnerId) = resolveCoupleAndPartner()
                val question = questionRepository.getQuestionById(questionId)
                val answer = localAnswerRepository.getAnswer(questionId)
                val partnerAnswer = if (coupleId != null && partnerId != null) {
                    runCatching {
                        firestoreAnswerDataSource.getAnswerForUser(
                            coupleId = coupleId,
                            userId = partnerId,
                            date = effectiveDate(answer)
                        )
                    }.onFailure { crashReporter.recordException(it) }.getOrNull()
                } else null

                val sealedPhase = computeSealedPhase(answer, partnerAnswer)
                val category = answer?.category ?: question?.category ?: ""

                _uiState.value = AnswerRevealUiState(
                    isLoading = false,
                    question = question,
                    answer = answer,
                    partnerAnswer = partnerAnswer,
                    coupleId = coupleId,
                    partnerId = partnerId,
                    sealedRevealPhase = sealedPhase,
                    followUpOptions = generateFollowUpOptions(answer, partnerAnswer, category)
                )

                // Live-watch the partner's answer doc so the reveal completes on this device
                // the moment they release their key — no manual refresh needed.
                if (coupleId != null && partnerId != null) {
                    observePartnerAnswerLive(coupleId, partnerId)
                }
            } catch (e: Exception) {
                crashReporter.recordException(e)
                _uiState.value = AnswerRevealUiState(
                    isLoading = false,
                    error = e.message ?: "Could not load this reveal."
                )
            }
        }
    }

    private fun computeSealedPhase(answer: LocalAnswer?, partnerAnswer: LocalAnswer?): SealedRevealPhase {
        if (answer?.schemaVersion != 3) return SealedRevealPhase.NONE
        if (answer.isRevealed) return SealedRevealPhase.REVEALED
        // isSealed == false means our key was already released; we are waiting for the partner's key.
        // Only check the local key store while the key is still pending (isSealed == true).
        if (!answer.isSealed) return SealedRevealPhase.WAITING_FOR_PARTNER
        val pendingKey = "daily_${answer.answerDate}_${questionId}"
        if (!pendingAnswerKeyStore.hasPendingKey(pendingKey)) return SealedRevealPhase.LOST_LOCAL_KEY
        return if (partnerAnswer != null) SealedRevealPhase.BOTH_ANSWERED else SealedRevealPhase.ANSWER_SEALED
    }

    private suspend fun resolveCoupleAndPartner(): Pair<String?, String?> {
        val userId = authRepository.currentUserId ?: return null to null
        val couple = runCatching { coupleRepository.getCoupleForUser(userId) }
            .onFailure { crashReporter.recordException(it) }
            .getOrNull()
            ?: return null to null
        val partnerId = couple.userIds.firstOrNull { it != userId }
        return couple.id to partnerId
    }

    private fun observeAnswer() {
        viewModelScope.launch {
            localAnswerRepository.observeAnswer(questionId).collect { answer ->
                _uiState.update { it.copy(answer = answer) }
            }
        }
    }

    /**
     * Watches the partner's answer doc live. When they submit (BOTH_ANSWERED appears) or
     * release their reveal key, the sealed phase is recomputed and — if our key is already
     * out — the partner's answer is decrypted on-device immediately. Decryption never leaves
     * the client; this only changes *when* it happens, not who can do it.
     */
    private fun observePartnerAnswerLive(coupleId: String, partnerId: String) {
        partnerObserveJob?.cancel()
        partnerObserveJob = viewModelScope.launch {
            val date = effectiveDate(_uiState.value.answer)
            firestoreAnswerDataSource.observeAnswerForUser(coupleId, partnerId, date).collect { partnerAnswer ->
                if (partnerAnswer == null) return@collect
                val current = _uiState.value
                // Once revealed, the partner answer in state is already decrypted — never
                // overwrite it with the sealed doc. Also leave an in-flight key release alone.
                if (current.sealedRevealPhase == SealedRevealPhase.REVEALED ||
                    current.sealedRevealPhase == SealedRevealPhase.RELEASING_KEY
                ) {
                    return@collect
                }
                val newPhase = computeSealedPhase(current.answer, partnerAnswer)
                _uiState.update { it.copy(partnerAnswer = partnerAnswer, sealedRevealPhase = newPhase) }
                // Our key is already released and the partner just released theirs → finish now.
                if (newPhase == SealedRevealPhase.WAITING_FOR_PARTNER) {
                    tryDecryptPartnerAnswer(coupleId, partnerId, _uiState.value)
                }
            }
        }
    }

    fun refreshPartnerAnswer() {
        val state = _uiState.value
        val coupleId = state.coupleId ?: return
        val partnerId = state.partnerId ?: return
        viewModelScope.launch {
            if (state.sealedRevealPhase == SealedRevealPhase.RELEASING_KEY) return@launch
            if (state.sealedRevealPhase == SealedRevealPhase.WAITING_FOR_PARTNER) {
                tryDecryptPartnerAnswer(coupleId, partnerId, state)
                return@launch
            }
            val partnerAnswer = runCatching {
                firestoreAnswerDataSource.getAnswerForUser(
                    coupleId = coupleId,
                    userId = partnerId,
                    date = effectiveDate(state.answer)
                )
            }.onFailure { crashReporter.recordException(it) }.getOrNull()
            partnerAnswer?.let { pa ->
                val newPhase = computeSealedPhase(state.answer, pa)
                _uiState.update { it.copy(partnerAnswer = pa, sealedRevealPhase = newPhase) }
            }
        }
    }

    fun revealAnswer() {
        val state = _uiState.value
        retentionAnalytics.track(RetentionEvent.RevealOpened(
            categoryId = state.answer?.category,
            coupleIdHash = state.coupleId?.coupleLoreHash()
        ))
        if (state.sealedRevealPhase == SealedRevealPhase.BOTH_ANSWERED) {
            performSealedReveal(state)
        } else {
            performLegacyReveal()
        }
    }

    private fun performSealedReveal(state: AnswerRevealUiState) {
        val coupleId = state.coupleId ?: return
        val partnerId = state.partnerId ?: return
        val userId = authRepository.currentUserId ?: return
        val date = effectiveDate(state.answer)

        _uiState.update { it.copy(sealedRevealPhase = SealedRevealPhase.RELEASING_KEY) }

        viewModelScope.launch {
            val released = runCatching {
                sealedRevealManager.releaseOwnKey(
                    coupleId = coupleId,
                    date = date,
                    questionId = questionId,
                    userId = userId,
                    partnerId = partnerId
                )
            }.onFailure { crashReporter.recordException(it) }.getOrDefault(false)

            if (!released) {
                _uiState.update { it.copy(sealedRevealPhase = SealedRevealPhase.LOST_LOCAL_KEY) }
                return@launch
            }

            // Update local answer to isSealed=false so cold restarts enter WAITING_FOR_PARTNER,
            // not LOST_LOCAL_KEY (the pending key was just removed from the store on release).
            val ownAnswer = _uiState.value.answer
            if (ownAnswer != null) {
                val updated = ownAnswer.copy(isSealed = false, updatedAt = System.currentTimeMillis())
                localAnswerRepository.saveAnswer(updated)
                _uiState.update { it.copy(answer = updated) }
            }
            // Mirror to Firestore so the partner's reveal screen can verify our key was released.
            runCatching {
                firestoreAnswerDataSource.markAnswerKeyReleased(coupleId, date, userId)
            }.onFailure { crashReporter.recordException(it) }

            tryDecryptPartnerAnswer(coupleId, partnerId, _uiState.value)
        }
    }

    private suspend fun tryDecryptPartnerAnswer(
        coupleId: String,
        partnerId: String,
        state: AnswerRevealUiState
    ) {
        val userId = authRepository.currentUserId ?: return
        val encryptedPayload = state.partnerAnswer?.encryptedPayload
        if (encryptedPayload == null) {
            _uiState.update { it.copy(sealedRevealPhase = SealedRevealPhase.WAITING_FOR_PARTNER) }
            return
        }

        val payload = runCatching {
            sealedRevealManager.decryptPartnerAnswer(
                coupleId = coupleId,
                date = effectiveDate(state.partnerAnswer),
                questionId = questionId,
                partnerId = partnerId,
                userId = userId,
                encryptedPayload = encryptedPayload
            )
        }.onFailure { crashReporter.recordException(it) }.getOrNull()

        if (payload == null) {
            _uiState.update { it.copy(sealedRevealPhase = SealedRevealPhase.WAITING_FOR_PARTNER) }
            return
        }

        // The sealed payload carries only option IDs; map them to labels via the question so the
        // reveal shows "A photo" rather than the raw "a_photo" id (own answer already has labels).
        val partnerOptionTexts = state.question
            ?.let { app.closer.ui.questions.selectedOptionTexts(it, payload.selectedOptionIds) }
            ?: emptyList()
        val decryptedPartnerAnswer = state.partnerAnswer?.copy(
            writtenText = payload.writtenText,
            selectedOptionIds = payload.selectedOptionIds,
            selectedOptionTexts = partnerOptionTexts,
            scaleValue = payload.scaleValue,
            isSealed = false
        )

        localAnswerRepository.markRevealed(questionId)
        val ownAnswer = localAnswerRepository.getAnswer(questionId)
        val category = ownAnswer?.category ?: state.question?.category ?: ""

        retentionAnalytics.track(RetentionEvent.RevealCompleted(
            categoryId = category,
            coupleIdHash = state.coupleId?.coupleLoreHash()
        ))

        _uiState.update {
            it.copy(
                answer = ownAnswer,
                partnerAnswer = decryptedPartnerAnswer,
                sealedRevealPhase = SealedRevealPhase.REVEALED,
                followUpOptions = generateFollowUpOptions(ownAnswer, decryptedPartnerAnswer, category)
            )
        }
    }

    private fun performLegacyReveal() {
        viewModelScope.launch {
            localAnswerRepository.markRevealed(questionId)
            val answer = localAnswerRepository.getAnswer(questionId)
            val partnerAnswer = _uiState.value.partnerAnswer
            val category = answer?.category ?: _uiState.value.question?.category ?: ""
            retentionAnalytics.track(RetentionEvent.RevealCompleted(
                categoryId = category,
                coupleIdHash = _uiState.value.coupleId?.coupleLoreHash()
            ))
            _uiState.update {
                it.copy(followUpOptions = generateFollowUpOptions(answer, partnerAnswer, category))
            }
        }
    }

    private fun generateFollowUpOptions(
        answer: LocalAnswer?,
        partnerAnswer: LocalAnswer?,
        category: String
    ): List<FollowUpOption> {
        if (answer == null || !answer.isRevealed) return emptyList()
        val bothRevealed = partnerAnswer?.isRevealed == true
        val options = mutableListOf<FollowUpOption>()
        options += FollowUpOption.DEEPER_FOLLOW_UP
        options += if (category.isNotBlank()) {
            FollowUpOption.ANOTHER_QUESTION
        } else if (bothRevealed) {
            FollowUpOption.DATE_IDEA
        } else {
            FollowUpOption.SAVE_MEMORY
        }
        return options.take(2)
    }

    fun onFollowUpSelected(option: FollowUpOption, onNavigate: (String) -> Unit) {
        val route = option.route
        if (route == null && option == FollowUpOption.ANOTHER_QUESTION) {
            val category = _uiState.value.answer?.category
                ?: _uiState.value.question?.category
                ?: return
            onNavigate(AppRoute.questionCategory(category))
            return
        }
        route?.let { onNavigate(it) } ?: showSnackbar("Coming soon")
    }

    fun saveCoupleLoré() {
        val state = _uiState.value
        val coupleId = state.coupleId ?: return
        val answer = state.answer ?: return
        if (state.loreSaved) return
        viewModelScope.launch {
            val modeTag = state.question?.tags?.firstOrNull { it.startsWith("mode_") }
            runCatching {
                firestoreAnswerDataSource.saveLoreEntry(
                    coupleId = coupleId,
                    questionId = answer.questionId,
                    questionText = state.question?.text ?: answer.questionText,
                    ownAnswer = answer.revealSummaryText(),
                    partnerAnswer = state.partnerAnswer?.revealSummaryText(),
                    modeTag = modeTag,
                    date = effectiveDate(answer)
                )
            }.onFailure { crashReporter.recordException(it) }
            _uiState.update { it.copy(loreSaved = true, snackbarMessage = "Saved to Couple Lore.") }
            retentionAnalytics.track(RetentionEvent.CoupleLoreSaved(
                categoryId = answer.category,
                coupleIdHash = coupleId.coupleLoreHash()
            ))
        }
    }

    fun onTinyActionSaved() {
        val state = _uiState.value
        retentionAnalytics.track(RetentionEvent.DailyTinyActionSaved(
            categoryId = state.answer?.category,
            coupleIdHash = state.coupleId?.coupleLoreHash()
        ))
    }

    fun showSnackbar(message: String) {
        _uiState.update { it.copy(snackbarMessage = message) }
    }

    fun clearSnackbar() {
        _uiState.update { it.copy(snackbarMessage = null) }
    }

    private fun String.coupleLoreHash(): String {
        return try {
            java.security.MessageDigest.getInstance("SHA-256")
                .digest(toByteArray())
                .take(8)
                .joinToString("") { "%02x".format(it) }
        } catch (e: Exception) { "hash_error" }
    }

    private fun LocalAnswer.revealSummaryText(): String = when (answerType) {
        "written" -> writtenText.orEmpty()
        "scale" -> "Chose ${scaleValue ?: "-"}"
        "single_choice", "multi_choice", "this_or_that" ->
            selectedOptionTexts.ifEmpty { selectedOptionIds }.joinToString()
        else -> writtenText ?: selectedOptionTexts.joinToString().ifBlank { "" }
    }

    private fun effectiveDate(answer: LocalAnswer?): String =
        answer?.answerDate?.takeIf { it.isNotBlank() }
            ?: FirestoreAnswerDataSource.todayLocalDateString()
}
