package app.closer.ui.pairing

import androidx.compose.foundation.background
import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.Spacer
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.foundation.layout.height
import androidx.compose.foundation.layout.heightIn
import androidx.compose.foundation.layout.navigationBarsPadding
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.layout.safeDrawingPadding
import androidx.compose.foundation.text.selection.SelectionContainer
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.Lock
import androidx.compose.material3.Button
import androidx.compose.material3.ButtonDefaults
import androidx.compose.material3.CircularProgressIndicator
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.OutlinedButton
import androidx.compose.material3.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.collectAsState
import androidx.compose.runtime.getValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.text.style.TextAlign
import androidx.compose.ui.unit.dp
import androidx.hilt.navigation.compose.hiltViewModel
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import app.closer.crypto.CoupleEncryptionManager
import app.closer.data.remote.CoupleAnswerMigrationDataSource
import app.closer.data.remote.FirestoreCoupleDataSource
import app.closer.domain.repository.AuthRepository
import app.closer.domain.repository.CoupleRepository
import app.closer.ui.components.BrandMessageRotator
import app.closer.ui.components.StatusGlyph
import app.closer.ui.settings.SettingsBackgroundBrush
import app.closer.ui.settings.SettingsInk
import app.closer.ui.settings.SettingsMuted
import app.closer.ui.settings.SettingsOnPrimary
import app.closer.ui.settings.SettingsPrimary
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 EncryptionUpgradeUiState(
    val isLoading: Boolean = true,
    val recoveryPhrase: String? = null,
    val complete: Boolean = false,
    val allPartnersComplete: Boolean = false,
    val error: String? = null
)

@HiltViewModel
class EncryptionUpgradeViewModel @Inject constructor(
    private val authRepository: AuthRepository,
    private val coupleRepository: CoupleRepository,
    private val coupleDataSource: FirestoreCoupleDataSource,
    private val migrationDataSource: CoupleAnswerMigrationDataSource,
    private val encryptionManager: CoupleEncryptionManager
) : ViewModel() {
    private val _uiState = MutableStateFlow(EncryptionUpgradeUiState())
    val uiState: StateFlow<EncryptionUpgradeUiState> = _uiState.asStateFlow()
    private var coupleId: String? = null

    init {
        upgrade()
    }

    fun upgrade() {
        if (_uiState.value.isLoading && coupleId != null) return
        _uiState.update { it.copy(isLoading = true, error = null) }
        viewModelScope.launch {
            runCatching {
                val userId = authRepository.currentUserId ?: error("Sign in again to secure your history.")
                var couple = coupleRepository.getCoupleForUser(userId)
                    ?: error("Your couple could not be loaded.")
                coupleId = couple.id

                if (couple.encryptionVersion >= CoupleEncryptionManager.STRICT_ENCRYPTION_VERSION) {
                    return@runCatching UpgradeResult(null, true)
                }

                var phrase = encryptionManager.pendingRecoveryPhrase(couple.id)
                if (couple.encryptionVersion == 0) {
                    val wrapped = if (phrase != null && encryptionManager.isUnlocked(couple.id)) {
                        encryptionManager.rewrapWithNewPhrase(couple.id, phrase).getOrThrow()
                    } else {
                        val setup = encryptionManager.setupLegacyCouple(couple.id)
                        phrase = setup.recoveryPhrase
                        setup.wrapped
                    }
                    val claimed = coupleDataSource.beginEncryptionMigration(couple.id, wrapped)
                    if (!claimed) {
                        encryptionManager.deleteKeyset(couple.id)
                        error("Your partner started the upgrade. Ask them for the recovery phrase, then unlock this device.")
                    }
                    couple = coupleRepository.getCoupleForUser(userId)
                        ?: error("Your couple could not be reloaded.")
                }

                if (!encryptionManager.isUnlocked(couple.id)) {
                    error("This device needs your shared recovery phrase before it can migrate answers.")
                }

                migrationDataSource.migrateUser(couple.id, userId)
                val allComplete = coupleDataSource.markEncryptionMigrationComplete(couple.id, userId)
                UpgradeResult(phrase, allComplete)
            }.onSuccess { result ->
                _uiState.update {
                    it.copy(
                        isLoading = false,
                        recoveryPhrase = result.phrase,
                        complete = true,
                        allPartnersComplete = result.allPartnersComplete
                    )
                }
            }.onFailure { error ->
                _uiState.update {
                    it.copy(isLoading = false, error = error.message ?: "The encryption upgrade failed. Try again.")
                }
            }
        }
    }

    fun acknowledgePhrase() {
        coupleId?.let(encryptionManager::acknowledgeRecoveryPhrase)
    }

    private data class UpgradeResult(val phrase: String?, val allPartnersComplete: Boolean)
}

@Composable
fun EncryptionUpgradeScreen(
    onComplete: () -> Unit,
    onRecoveryNeeded: () -> Unit,
    viewModel: EncryptionUpgradeViewModel = hiltViewModel()
) {
    val state by viewModel.uiState.collectAsState()

    Column(
        modifier = Modifier
            .fillMaxSize()
            .background(SettingsBackgroundBrush)
            .safeDrawingPadding()
            .navigationBarsPadding()
            .padding(horizontal = 28.dp, vertical = 32.dp),
        horizontalAlignment = Alignment.CenterHorizontally,
        verticalArrangement = Arrangement.Center
    ) {
        StatusGlyph(
            icon = Icons.Filled.Lock,
            tint = SettingsPrimary,
            container = SettingsPrimary.copy(alpha = 0.14f)
        )
        Spacer(Modifier.height(20.dp))
        Text(
            text = if (state.complete) "Your answers are secured" else "Securing your history",
            style = MaterialTheme.typography.headlineSmall,
            color = SettingsInk,
            fontWeight = FontWeight.SemiBold,
            textAlign = TextAlign.Center
        )
        Spacer(Modifier.height(10.dp))

        when {
            state.isLoading -> {
                Text(
                    "This one-time upgrade encrypts answer content on this device before replacing its cloud copy.",
                    style = MaterialTheme.typography.bodyMedium,
                    color = SettingsMuted,
                    textAlign = TextAlign.Center
                )
                Spacer(Modifier.height(24.dp))
                CircularProgressIndicator(color = SettingsPrimary)
                Spacer(Modifier.height(16.dp))
                BrandMessageRotator(style = MaterialTheme.typography.bodySmall)
            }
            state.error != null -> {
                val error = state.error.orEmpty()
                Text(
                    error,
                    style = MaterialTheme.typography.bodyMedium,
                    color = MaterialTheme.colorScheme.error,
                    textAlign = TextAlign.Center
                )
                Spacer(Modifier.height(20.dp))
                Button(onClick = viewModel::upgrade) { Text("Try again") }
                if (error.contains("recovery phrase", ignoreCase = true)) {
                    OutlinedButton(onClick = onRecoveryNeeded) { Text("Enter recovery phrase") }
                }
            }
            state.recoveryPhrase != null -> {
                val recoveryPhrase = state.recoveryPhrase.orEmpty()
                Text(
                    "Save this phrase somewhere private and share it directly with your partner. Closer does not store the phrase and cannot recover it.",
                    style = MaterialTheme.typography.bodyMedium,
                    color = SettingsMuted,
                    textAlign = TextAlign.Center
                )
                Spacer(Modifier.height(20.dp))
                SelectionContainer {
                    Text(
                        recoveryPhrase,
                        modifier = Modifier.fillMaxWidth().padding(16.dp),
                        style = MaterialTheme.typography.titleMedium,
                        color = SettingsInk,
                        textAlign = TextAlign.Center,
                        fontWeight = FontWeight.SemiBold
                    )
                }
                Spacer(Modifier.height(20.dp))
                Button(
                    onClick = { viewModel.acknowledgePhrase(); onComplete() },
                    modifier = Modifier.fillMaxWidth().heightIn(min = 54.dp),
                    colors = ButtonDefaults.buttonColors(
                        containerColor = SettingsPrimary,
                        contentColor = SettingsOnPrimary
                    )
                ) { Text("I've saved and shared it") }
            }
            else -> {
                Text(
                    if (state.allPartnersComplete)
                        "Both sides have migrated. New and historical answer content now uses strict end-to-end encryption."
                    else
                        "This device is ready. Your partner will finish the upgrade after unlocking with the shared recovery phrase.",
                    style = MaterialTheme.typography.bodyMedium,
                    color = SettingsMuted,
                    textAlign = TextAlign.Center
                )
                Spacer(Modifier.height(24.dp))
                Button(onClick = onComplete, modifier = Modifier.fillMaxWidth()) { Text("Continue") }
            }
        }
    }
}
