package com.peachpasswords.ui.pairing

import android.content.Context
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import com.peachpasswords.license.LicenseRepository
import com.peachpasswords.sync.AdoptionReceiverDeps
import com.peachpasswords.sync.AdoptionSourceDeps
import com.peachpasswords.sync.BackgroundSyncPreference
import com.peachpasswords.sync.DeviceType
import com.peachpasswords.sync.AdoptionFrame
import com.peachpasswords.sync.LocalDeviceIdentity
import com.peachpasswords.sync.PairedDevice
import com.peachpasswords.sync.PairedDeviceStore
import com.peachpasswords.sync.SignedDeviceManifest
import com.peachpasswords.sync.VaultSyncAdapter
import com.peachpasswords.sync.PeachSyncService
import com.peachpasswords.sync.RelayAdoptionReceiver
import com.peachpasswords.sync.RelaySettingsStore
import com.peachpasswords.sync.RelayDecision
import com.peachpasswords.sync.VaultStatus
import com.peachpasswords.sync.buildLocalIdentity
import com.peachpasswords.sync.computeVaultId
import com.peachpasswords.sync.deriveLocalSyncKeypair
import com.peachpasswords.sync.getDeviceName
import com.peachpasswords.sync.getOrCreateDeviceId
import com.peachpasswords.sync.hexToPubkey
import com.peachpasswords.sync.parsePairingInvitation
import com.peachpasswords.ui.security.scorePassword
import com.peachpasswords.vault.AdoptVaultResult
import com.peachpasswords.vault.VaultCrypto
import com.peachpasswords.vault.VaultRepository
import dagger.hilt.android.lifecycle.HiltViewModel
import dagger.hilt.android.qualifiers.ApplicationContext
import java.util.UUID
import javax.inject.Inject
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.Job
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.launch

sealed interface AdoptVaultUiState {
    data object Password : AdoptVaultUiState
    data object Pair : AdoptVaultUiState
    data object Connecting : AdoptVaultUiState
    /**
     * Source-initiator flow exposes [isSource] so the screen can phrase the
     * confirmation as "send vault" instead of "adopt/join vault". The SAS is
     * identical to the receiver flow; only the payload direction differs.
     */
    data class ConfirmSas(
        val sas: String,
        val isSource: Boolean = false,
        val peerName: String? = null,
    ) : AdoptVaultUiState
    data object Adopting : AdoptVaultUiState
    data class Complete(val sourceName: String) : AdoptVaultUiState
    data class Failed(val reason: String) : AdoptVaultUiState
}

@HiltViewModel
class AdoptVaultViewModel @Inject constructor(
    private val receiver: RelayAdoptionReceiver,
    private val vault: VaultRepository,
    private val devices: PairedDeviceStore,
    private val relaySettings: RelaySettingsStore,
    private val backgroundSync: BackgroundSyncPreference,
    private val license: LicenseRepository,
    @ApplicationContext private val appContext: Context,
) : ViewModel() {
    /**
     * True when this device already has a vault. A populated device may refresh
     * a pairing only for that exact vault identity; it may not adopt a different
     * root key and silently merge two record sets.
     */
    val joining: Boolean = vault.hasVault()

    private val _state = MutableStateFlow<AdoptVaultUiState>(
        initialAdoptVaultState(joining),
    )
    val state: StateFlow<AdoptVaultUiState> = _state.asStateFlow()

    private var password: String? = null
    private var confirmation: CompletableDeferred<Boolean>? = null
    private var adoptionJob: Job? = null
    private var protocolError: String? = null

    fun acceptPassword(value: String, repeated: String): Boolean {
        if (value.length < 12 || value != repeated || scorePassword(value) < 2) {
            return false
        }
        password = value
        _state.value = AdoptVaultUiState.Pair
        return true
    }

    fun adopt(invitation: String) {
        if (_state.value !is AdoptVaultUiState.Pair) return
        val parsed = parsePairingInvitation(invitation)
        if (parsed == null) {
            _state.value = AdoptVaultUiState.Failed("invalid_pairing_invitation")
            clearPassword()
            return
        }
        when (parsed.vaultStatus) {
            VaultStatus.POPULATED -> startReceiverFlow(invitation)
            VaultStatus.EMPTY -> startSourceFlow(invitation)
        }
    }

    private fun startReceiverFlow(invitation: String) {
        _state.value = AdoptVaultUiState.Connecting
        protocolError = null
        adoptionJob = viewModelScope.launch {
            try {
                val adopted = receiver.adopt(
                    encodedInvitation = invitation,
                    deps = receiverDeps(),
                    confirm = { sas ->
                        val pending = CompletableDeferred<Boolean>()
                        confirmation = pending
                        _state.value = AdoptVaultUiState.ConfirmSas(
                            sas = sas,
                            isSource = false,
                        )
                        pending.await()
                    },
                )
                // The pair carries the source's relay, but the Sync card reads
                // this device's own setting — left alone it would say "Public
                // relay" on a phone dialling a private one. Peach's own URL
                // still lands on PUBLIC, so joining over the public relay is
                // not relabelled private.
                adopted.relayUrl?.let(relaySettings::setRelayUrl)
                _state.value = AdoptVaultUiState.Complete(
                    adopted.deviceName?.takeIf(String::isNotBlank) ?: "Your other device",
                )
            } catch (error: CancellationException) {
                throw error
            } catch (error: Throwable) {
                _state.value = AdoptVaultUiState.Failed(
                    protocolError ?: error.message ?: "adoption_failed",
                )
            } finally {
                confirmation = null
                clearPassword()
            }
        }
    }

    /**
     * Source-initiator flow: the scanning invitation came from a vault-less
     * device, so the phone dials the relay, completes Noise, and sends the
     * root key over. Uses the same UI states as the receiver flow so the
     * surface area stays the same; only [AdoptVaultUiState.ConfirmSas.isSource]
     * distinguishes the role for the confirmation copy.
     */
    private fun startSourceFlow(invitation: String) {
        val master = vault.syncMasterKey()
        if (master == null) {
            _state.value = AdoptVaultUiState.Failed("vault_locked")
            clearPassword()
            return
        }
        when (val decision = relaySettings.relayInForce()) {
            is RelayDecision.Refuse -> {
                VaultCrypto.zero(master)
                _state.value = AdoptVaultUiState.Failed(decision.reason.wire())
                clearPassword()
                return
            }
            is RelayDecision.Dial -> Unit
        }
        _state.value = AdoptVaultUiState.Connecting
        protocolError = null

        val keys = deriveLocalSyncKeypair(master, getOrCreateDeviceId(appContext))
        val identity = buildLocalIdentity(appContext, master)
        VaultCrypto.zero(master)
        adoptionJob = viewModelScope.launch {
            try {
                val paired = receiver.grantFromInvitation(
                    encodedInvitation = invitation,
                    sourceDeps = sourceDeps(identity),
                    confirm = { sas, request ->
                        val pending = CompletableDeferred<Boolean>()
                        confirmation = pending
                        _state.value = AdoptVaultUiState.ConfirmSas(
                            sas = sas,
                            isSource = true,
                            peerName = request.deviceName.ifBlank { "New device" },
                        )
                        pending.await()
                    },
                )
                paired.relayUrl?.let(relaySettings::setRelayUrl)
                _state.value = AdoptVaultUiState.Complete(
                    paired.deviceName?.takeIf(String::isNotBlank) ?: "Your other device",
                )
            } catch (error: CancellationException) {
                throw error
            } catch (error: Throwable) {
                _state.value = AdoptVaultUiState.Failed(
                    protocolError ?: error.message ?: "adoption_failed",
                )
            } finally {
                VaultCrypto.zero(keys.secretKey)
                VaultCrypto.zero(keys.publicKey)
                confirmation = null
                clearPassword()
            }
        }
    }

    fun confirmSas(matches: Boolean) {
        val pending = confirmation ?: return
        confirmation = null
        if (matches) _state.value = AdoptVaultUiState.Adopting
        pending.complete(matches)
    }

    fun restart() {
        adoptionJob?.cancel()
        adoptionJob = null
        confirmation?.complete(false)
        confirmation = null
        clearPassword()
        _state.value = initialAdoptVaultState(joining)
    }

    private fun receiverDeps(): AdoptionReceiverDeps =
        object : AdoptionReceiverDeps {
            override val deviceName = getDeviceName(appContext)
            override val deviceType = DeviceType.ANDROID
            override fun generateDeviceId(): String = UUID.randomUUID().toString()

            // RelayAdoptionReceiver replaces this with the SAS bound to its
            // completed Noise transcript.
            override suspend fun confirmSas(): Boolean = false

            override suspend fun validateSource(source: AdoptionFrame.Source) {
                // First contact on a populated vault is allowed: a second
                // browser, a tablet, another phone — any device that holds the
                // same root key (the same recovery words) is the same vault.
                // The vault-id check in createVaultFromRootKey is the real
                // security boundary; this callback now only documents intent
                // for callers that want to layer their own policy on top.
                if (!joining) return
            }

            override suspend fun createVaultFromRootKey(rootKeyHex: String): String {
                // Same protocol seam, two safe outcomes: a fresh device builds
                // around the granted key; a populated device may only refresh a
                // pairing when it already has that exact vault identity.
                val result = if (joining) {
                    val incomingRoot = hexToPubkey(rootKeyHex)
                    val incomingVaultId = try {
                        computeVaultId(incomingRoot)
                    } finally {
                        VaultCrypto.zero(incomingRoot)
                    }
                    AdoptVaultResult.Success(
                        requireSameVaultForPopulatedRePair(
                            localVaultId = vault.syncVaultId(),
                            incomingVaultId = incomingVaultId,
                        ),
                    )
                } else {
                    val localPassword = password
                        ?: throw IllegalStateException("adoption_password_unavailable")
                    vault.createVaultFromRootKey(localPassword, rootKeyHex)
                }
                return when (result) {
                    is AdoptVaultResult.Success -> result.vaultId
                    is AdoptVaultResult.Error -> throw IllegalStateException(result.reason)
                }
            }

            override suspend fun deriveLocalIdentity(): LocalDeviceIdentity {
                val master = vault.syncMasterKey()
                    ?: throw IllegalStateException("adopted_vault_unavailable")
                return try {
                    buildLocalIdentity(appContext, master)
                } finally {
                    VaultCrypto.zero(master)
                }
            }

            override suspend fun clearPairedDeviceRevocation(deviceId: String) {
                devices.clearRevocation(deviceId)
            }

            override suspend fun upsertPairedDevice(device: PairedDevice) {
                devices.upsert(device)
            }

            override suspend fun applyLicense(licenseKey: String) {
                if (!license.validateLicense(licenseKey)) {
                    throw IllegalStateException("license_transfer_failed")
                }
                license.storeLicense(licenseKey)
                if (!license.isProActive()) {
                    throw IllegalStateException("license_transfer_failed")
                }
            }

            override suspend fun license(): String? =
                license.getLicenseKey()?.takeIf(license::validateLicense)

            override fun onError(error: String) {
                protocolError = error
            }

            override suspend fun applySnapshot(entries: Map<String, String>, tombstones: Map<String, Long>) {
                val adapter = VaultSyncAdapter(
                    repo = vault,
                    ourDeviceId = getOrCreateDeviceId(appContext),
                    pairedDevices = devices,
                )
                adapter.applyEntries(entries)
                adapter.applyTombstones(tombstones)
            }

            override suspend fun markSynced(deviceId: String) {
                devices.updateLastSync(deviceId)
            }
        }

    private fun sourceDeps(identity: LocalDeviceIdentity): AdoptionSourceDeps =
        object : AdoptionSourceDeps {
            override val sourceIdentity: LocalDeviceIdentity = identity

            override suspend fun rootKeyHex(): String {
                val root = vault.syncMasterKey()
                    ?: throw IllegalStateException("vault_locked")
                return try {
                    root.joinToString("") { "%02x".format(it) }
                } finally {
                    VaultCrypto.zero(root)
                }
            }

            override suspend fun vaultId(): String =
                vault.syncVaultId() ?: throw IllegalStateException("vault_locked")

            override suspend fun license(): String? =
                license.getLicenseKey()?.takeIf(license::validateLicense)

            override suspend fun applyLicense(licenseKey: String) {
                if (!license.validateLicense(licenseKey)) {
                    throw IllegalStateException("license_transfer_failed")
                }
                license.storeLicense(licenseKey)
                if (!license.isProActive()) {
                    throw IllegalStateException("license_transfer_failed")
                }
            }

            override suspend fun onConsent(request: AdoptionFrame.Request): Boolean = true

            override suspend fun upsertPairedDevice(device: PairedDevice) {
                devices.clearRevocation(device.deviceId)
                devices.upsert(device)
            }

            override suspend fun vaultSnapshot(): Pair<Map<String, String>, Map<String, Long>>? {
                val adapter = VaultSyncAdapter(
                    repo = vault,
                    ourDeviceId = getOrCreateDeviceId(appContext),
                    pairedDevices = devices,
                )
                val entryManifest = adapter.buildEntryManifest()
                val entryIds = entryManifest.keys.toList()
                val entries = if (entryIds.isNotEmpty()) adapter.getEncryptedEntries(entryIds) else emptyMap()
                val tombstones = adapter.buildTombstoneManifest()
                return entries to tombstones
            }

            override suspend fun outgoingDeviceManifest(): SignedDeviceManifest? {
                val adapter = VaultSyncAdapter(
                    repo = vault,
                    ourDeviceId = getOrCreateDeviceId(appContext),
                    pairedDevices = devices,
                )
                return adapter.buildOutgoingDeviceManifest()
            }

            override suspend fun markSynced(deviceId: String) {
                devices.updateLastSync(deviceId)
            }
        }

    fun setBackgroundSync(enabled: Boolean) {
        backgroundSync.setEnabled(enabled)
        if (enabled) PeachSyncService.start(appContext) else PeachSyncService.stop(appContext)
    }

    private fun clearPassword() {
        password = null
    }

    override fun onCleared() {
        adoptionJob?.cancel()
        confirmation?.complete(false)
        clearPassword()
        super.onCleared()
    }
}

internal fun initialAdoptVaultState(hasVault: Boolean): AdoptVaultUiState =
    if (hasVault) AdoptVaultUiState.Pair else AdoptVaultUiState.Password

/**
 * Pairing is not a merge tool. A populated destination may re-pair only when
 * both sides already derive the same vault id. Different ids must stop before
 * recovery words, entries, revocations, or paired-device state are changed.
 */
internal fun requireSameVaultForPopulatedRePair(
    localVaultId: String?,
    incomingVaultId: String,
): String {
    if (localVaultId == null) throw IllegalStateException("vault_locked")
    if (localVaultId != incomingVaultId) {
        throw IllegalStateException("JOIN_POPULATED_VAULT_BLOCKED")
    }
    return localVaultId
}