Closer/app/src/main/java/app/closer/ui/home/HomeViewModel.kt

package app.closer.ui.home

import android.util.Log
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import app.closer.crypto.CoupleEncryptionManager
import app.closer.crypto.EncryptionStatus
import app.closer.domain.model.LocalAnswer
import app.closer.domain.model.Question
import app.closer.domain.model.QuestionCategory
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 app.closer.domain.repository.UserRepository
import com.google.firebase.firestore.FirebaseFirestore
import dagger.hilt.android.lifecycle.HiltViewModel
import javax.inject.Inject
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch

data class HomeCategorySummary(
    val category: QuestionCategory,
    val questionCount: Int
)

data class HomeAnswerStats(
    val total: Int = 0,
    val revealed: Int = 0,
    val private: Int = 0,
    val latest: LocalAnswer? = null,
    val answeredQuestionIds: Set<String> = emptySet()
)

enum class HomeActionTarget {
    InvitePartner,
    DailyQuestion,
    AnswerHistory,
    QuestionPacks,
    Settings
}

enum class HomeActionTone {
    Invite,
    Daily,
    Reflection,
    Ritual,
    Starter,
    Pack,
    Utility
}

data class HomeAction(
    val eyebrow: String,
    val title: String,
    val body: String,
    val cta: String,
    val target: HomeActionTarget,
    val tone: HomeActionTone,
    val metric: String? = null,
    val categoryId: String? = null
)

enum class DailyQuestionState {
    UNANSWERED,
    USER_ANSWERED_PARTNER_PENDING,
    PARTNER_ANSWERED_USER_PENDING,
    BOTH_ANSWERED,
    REVEALED
}

data class HomeUiState(
    val isLoading: Boolean = true,
    val error: String? = null,
    val dailyQuestion: Question? = null,
    val categories: List<HomeCategorySummary> = emptyList(),
    val answerStats: HomeAnswerStats = HomeAnswerStats(),
    val partnerName: String? = null,
    val streakCount: Int = 0,
    val isPaired: Boolean = false,
    val primaryAction: HomeAction? = null,
    val secondaryActions: List<HomeAction> = emptyList(),
    val partnerLeftEvent: Boolean = false,
    val needsRecovery: Boolean = false,
    val needsEncryptionUpgrade: Boolean = false,
    val dailyQuestionState: DailyQuestionState = DailyQuestionState.UNANSWERED,
    val hasPartnerAnsweredToday: Boolean = false,
    val partnerAnsweredQuestionId: String? = null,
    val hasRevealedToday: Boolean = false
)

@HiltViewModel
class HomeViewModel @Inject constructor(
    private val questionRepository: QuestionRepository,
    private val localAnswerRepository: LocalAnswerRepository,
    private val authRepository: AuthRepository,
    private val coupleRepository: CoupleRepository,
    private val userRepository: UserRepository,
    private val encryptionManager: CoupleEncryptionManager,
    private val db: FirebaseFirestore
) : ViewModel() {

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

    private var coupleStateListener: com.google.firebase.firestore.ListenerRegistration? = null
    private var partnerAnswerListener: com.google.firebase.firestore.ListenerRegistration? = null

    init {
        loadHome()
        observeAnswers()
        observeCoupleState()
    }

    override fun onCleared() {
        super.onCleared()
        coupleStateListener?.remove()
        coupleStateListener = null
    }

    fun loadHome() {
        viewModelScope.launch {
            _uiState.update { it.copy(isLoading = true, error = null) }
            try {
                val dailyQuestion = questionRepository.getDailyQuestion()
                val categories = questionRepository.getCategories()
                    .take(6)
                    .map { category ->
                        HomeCategorySummary(
                            category = category,
                            questionCount = questionRepository.getQuestionCountByCategory(category.id)
                        )
                    }
                val uid = authRepository.currentUserId
                val couple = uid?.let { runCatching { coupleRepository.getCoupleForUser(it) }.getOrNull() }
                val partnerName = couple?.userIds?.firstOrNull { it != uid }?.let { partnerId ->
                    runCatching { userRepository.getUser(partnerId)?.displayName }
                        .onFailure { Log.w(TAG, "Could not load partner display name", it) }
                        .getOrNull()
                }
                val encryptionStatus = couple?.let(encryptionManager::checkStatus)
                val needsRecovery = encryptionStatus == EncryptionStatus.NEEDS_RECOVERY
                val needsEncryptionUpgrade = when (encryptionStatus) {
                    EncryptionStatus.NEEDS_ENCRYPTION_UPGRADE -> true
                    EncryptionStatus.NEEDS_CONTENT_MIGRATION ->
                        couple.encryptionMigrationUsers[uid] != true
                    else -> false
                }
                _uiState.update { current ->
                    current.copy(
                        isLoading = false,
                        dailyQuestion = dailyQuestion,
                        categories = categories,
                        partnerName = partnerName,
                        streakCount = couple?.streakCount ?: 0,
                        isPaired = couple != null,
                        partnerLeftEvent = false,
                        needsRecovery = needsRecovery,
                        needsEncryptionUpgrade = needsEncryptionUpgrade
                    ).refreshDailyQuestionState().withHomeActions()
                }
                observePartnerAnswer(couple?.id, couple?.userIds, dailyQuestion?.id)
            } catch (e: Exception) {
                _uiState.update {
                    it.copy(
                        isLoading = false,
                        error = e.message ?: "Couldn't load your space right now."
                    ).withHomeActions()
                }
            }
        }
    }

    private fun observeCoupleState() {
        val uid = authRepository.currentUserId ?: return
        coupleStateListener?.remove()
        coupleStateListener = db.collection("users").document(uid)
            .addSnapshotListener(com.google.firebase.firestore.MetadataChanges.INCLUDE) { snapshot, error ->
                if (error != null || snapshot == null || !snapshot.exists()) {
                    return@addSnapshotListener
                }
                val newCoupleId = snapshot.getString("coupleId")
                val oldIsPaired = _uiState.value.isPaired
                val oldPartnerName = _uiState.value.partnerName
                val newIsPaired = !newCoupleId.isNullOrBlank()

                if (newIsPaired != oldIsPaired) {
                    // Capture partner name before reload so the banner can reference it.
                    val showPartnerLeftBanner = oldIsPaired && !newIsPaired
                    if (showPartnerLeftBanner) {
                        _uiState.update {
                            it.copy(
                                isPaired = false,
                                partnerName = null,
                                partnerLeftEvent = true
                            ).withHomeActions()
                        }
                    }
                    loadHome()
                }
            }
    }

    /**
     * Consumes the partner-left banner event. Call from the UI after showing it.
     */
    fun consumePartnerLeftEvent() {
        _uiState.update { it.copy(partnerLeftEvent = false) }
    }

    /** Called after the recovery flow completes so the banner goes away. */
    fun onRecoveryCompleted() {
        _uiState.update { it.copy(needsRecovery = false) }
        loadHome()
    }

    private fun observeAnswers() {
        viewModelScope.launch {
            localAnswerRepository.observeAnswers().collect { answers ->
                val sorted = answers.sortedByDescending { it.updatedAt }
                _uiState.update {
                    it.copy(
                        answerStats = HomeAnswerStats(
                            total = answers.size,
                            revealed = answers.count { answer -> answer.isRevealed },
                            private = answers.count { answer -> !answer.isRevealed },
                            latest = sorted.firstOrNull(),
                            answeredQuestionIds = answers.map { answer -> answer.questionId }.toSet()
                        )
                    ).refreshDailyQuestionState().withHomeActions()
                }
            }
        }
    }

    private fun observePartnerAnswer(
        coupleId: String?,
        coupleUserIds: List<String>?,
        dailyQuestionId: String?
    ) {
        partnerAnswerListener?.remove()
        partnerAnswerListener = null
        val cId = coupleId ?: return
        val qId = dailyQuestionId ?: return
        val uid = authRepository.currentUserId ?: return
        val partnerId = coupleUserIds?.firstOrNull { it != uid } ?: return

        val today = FirestoreAnswerDataSource.todayLocalDateString()
        partnerAnswerListener = db.collection("couples")
            .document(cId)
            .collection("daily_question")
            .document(today)
            .collection("answers")
            .document(partnerId)
            .addSnapshotListener(com.google.firebase.firestore.MetadataChanges.INCLUDE) { snapshot, error ->
                if (error != null) {
                    Log.w(TAG, "Partner answer listener failed", error)
                    return@addSnapshotListener
                }
                val hasPartnerAnswer = snapshot?.exists() == true
                _uiState.update {
                    it.copy(
                        hasPartnerAnsweredToday = hasPartnerAnswer,
                        partnerAnsweredQuestionId = if (hasPartnerAnswer) qId else null
                    ).refreshDailyQuestionState().withHomeActions()
                }
            }
    }

    private fun HomeUiState.refreshDailyQuestionState(): HomeUiState {
        val questionId = dailyQuestion?.id
        val userAnswered = questionId != null && questionId in answerStats.answeredQuestionIds
        val userRevealed = questionId != null && answerStats.latest?.let { latest ->
            latest.questionId == questionId && latest.isRevealed
        } == true

        val state = when {
            questionId == null -> DailyQuestionState.UNANSWERED
            userRevealed -> DailyQuestionState.REVEALED
            userAnswered && hasPartnerAnsweredToday -> DailyQuestionState.BOTH_ANSWERED
            userAnswered -> DailyQuestionState.USER_ANSWERED_PARTNER_PENDING
            hasPartnerAnsweredToday -> DailyQuestionState.PARTNER_ANSWERED_USER_PENDING
            else -> DailyQuestionState.UNANSWERED
        }
        return copy(
            dailyQuestionState = state,
            hasRevealedToday = userRevealed
        )
    }

    private fun HomeUiState.withHomeActions(): HomeUiState {
        if (isLoading || error != null) {
            return copy(primaryAction = null, secondaryActions = emptyList())
        }

        val primary = buildPrimaryAction()
        return copy(
            primaryAction = primary,
            secondaryActions = buildSecondaryActions(primary)
        )
    }

    private fun HomeUiState.buildPrimaryAction(): HomeAction {
        val dailyQuestionId = dailyQuestion?.id
        val userAnswered = dailyQuestionId != null && dailyQuestionId in answerStats.answeredQuestionIds
        val userRevealed = dailyQuestionId != null && answerStats.latest?.let { latest ->
            latest.questionId == dailyQuestionId && latest.isRevealed
        } == true

        return when {
            !isPaired -> HomeAction(
                eyebrow = "Next best action",
                title = "Invite your partner into tonight.",
                body = "The app works best as a shared ritual. Send a private invite and make the next prompt something you can both answer.",
                cta = "Invite partner",
                target = HomeActionTarget.InvitePartner,
                tone = HomeActionTone.Invite
            )

            userRevealed -> HomeAction(
                eyebrow = "Tonight's prompt",
                title = "You opened a conversation tonight.",
                body = dailyQuestion?.text ?: "You revealed an answer together. What comes next is up to both of you.",
                cta = "Try a follow-up",
                target = HomeActionTarget.DailyQuestion,
                tone = HomeActionTone.Daily,
                metric = dailyQuestion?.category?.takeIf { category -> category.isNotBlank() }?.toHomeLabel()
            )

            dailyQuestionState == DailyQuestionState.BOTH_ANSWERED -> HomeAction(
                eyebrow = "Tonight's prompt",
                title = "Reveal is ready.",
                body = "Both of you answered. Open it together when you are both in the right headspace.",
                cta = "Reveal together",
                target = HomeActionTarget.DailyQuestion,
                tone = HomeActionTone.Daily,
                metric = dailyQuestion?.category?.takeIf { category -> category.isNotBlank() }?.toHomeLabel()
            )

            dailyQuestionState == DailyQuestionState.USER_ANSWERED_PARTNER_PENDING -> HomeAction(
                eyebrow = "Tonight's prompt",
                title = "You showed up tonight. Waiting for your partner.",
                body = "Your answer is private until they answer too. No pressure — the reveal waits for both of you.",
                cta = "Send a gentle reminder",
                target = HomeActionTarget.DailyQuestion,
                tone = HomeActionTone.Daily,
                metric = dailyQuestion?.category?.takeIf { category -> category.isNotBlank() }?.toHomeLabel()
            )

            dailyQuestionState == DailyQuestionState.PARTNER_ANSWERED_USER_PENDING -> HomeAction(
                eyebrow = "Tonight's prompt",
                title = "Your partner answered. Your turn.",
                body = "Answer to unlock the reveal. Your response stays private until you are ready.",
                cta = "Answer to unlock reveal",
                target = HomeActionTarget.DailyQuestion,
                tone = HomeActionTone.Daily,
                metric = dailyQuestion?.category?.takeIf { category -> category.isNotBlank() }?.toHomeLabel()
            )

            userAnswered -> HomeAction(
                eyebrow = "Tonight's prompt",
                title = dailyQuestion?.text ?: "Answer tonight's question.",
                body = "Start with one honest answer. You can keep it private or reveal it when the moment feels right.",
                cta = "Answer now",
                target = HomeActionTarget.DailyQuestion,
                tone = HomeActionTone.Daily,
                metric = dailyQuestion?.category?.takeIf { category -> category.isNotBlank() }?.toHomeLabel()
            )

            answerStats.private > 0 -> HomeAction(
                eyebrow = "Saved privately",
                title = "You have ${answerStats.private} reflection${if (answerStats.private == 1) "" else "s"} waiting.",
                body = "Review what you saved and choose whether tonight is the right time to open one up.",
                cta = "Review reflections",
                target = HomeActionTarget.AnswerHistory,
                tone = HomeActionTone.Reflection,
                metric = "${answerStats.revealed} revealed"
            )

            streakCount > 0 -> HomeAction(
                eyebrow = "Shared ritual",
                title = "$streakCount night${if (streakCount == 1) "" else "s"} showing up.",
                body = "Keep it light: answer one prompt, revisit a saved reflection, or choose a pack that fits tonight.",
                cta = "Keep going",
                target = HomeActionTarget.DailyQuestion,
                tone = HomeActionTone.Ritual,
                metric = "${answerStats.total} saved"
            )

            else -> HomeAction(
                eyebrow = "Gentle start",
                title = "Tonight's question is ready.",
                body = "${dailyQuestion?.text ?: "A small prompt is enough. Build the habit around attention, not pressure."}",
                cta = "Answer privately",
                target = HomeActionTarget.DailyQuestion,
                tone = HomeActionTone.Starter,
                metric = dailyQuestion?.category?.takeIf { category -> category.isNotBlank() }?.toHomeLabel()
            )
        }
    }

    private fun HomeUiState.buildSecondaryActions(primary: HomeAction): List<HomeAction> {
        val actions = mutableListOf<HomeAction>()

        answerStats.latest?.let { latest ->
            if (primary.target != HomeActionTarget.AnswerHistory) {
                actions += HomeAction(
                    eyebrow = if (latest.isRevealed) "Revealed" else "Private",
                    title = "Return to your latest reflection.",
                    body = latest.questionText,
                    cta = "Open history",
                    target = HomeActionTarget.AnswerHistory,
                    tone = HomeActionTone.Reflection
                )
            }
        }

        categories.firstOrNull()?.let { category ->
            actions += HomeAction(
                eyebrow = "Suggested pack",
                title = category.category.displayName.ifBlank { "Question pack" },
                body = "${category.questionCount} prompts for when you want a different doorway into the conversation.",
                cta = "Open pack",
                target = HomeActionTarget.QuestionPacks,
                tone = HomeActionTone.Pack,
                categoryId = category.category.id
            )
        }

        actions += HomeAction(
            eyebrow = "Tune the ritual",
            title = "Adjust your space.",
            body = "Manage reminders, partner state, privacy, and account details when you need to.",
            cta = "Settings",
            target = HomeActionTarget.Settings,
            tone = HomeActionTone.Utility
        )

        return actions.take(3)
    }

    private fun String.toHomeLabel(): String =
        split("_", "-")
            .filter { part -> part.isNotBlank() }
            .joinToString(" ") { part -> part.replaceFirstChar { it.uppercaseChar() } }

    companion object {
        private const val TAG = "HomeViewModel"
    }
}