package com.peachpasswords.ui.pairing

import android.content.Context
import android.util.Log
import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import com.peachpasswords.license.LicenseRepository
import com.peachpasswords.sync.AdoptionPairState
import com.peachpasswords.sync.AdoptionReceiverDeps
import com.peachpasswords.sync.AdoptionReceiverV2Deps
import com.peachpasswords.sync.AdoptionSourceDeps
import com.peachpasswords.sync.AdoptionSourceV2Deps
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.PEER_PROTOCOL_OUTDATED
import com.peachpasswords.sync.PendingVouchedDevice
import com.peachpasswords.sync.PendingVouchedDeviceStore
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.getDeviceName
import com.peachpasswords.sync.hexToPubkey
import com.peachpasswords.sync.loadOrCreateLocalIdentity
import com.peachpasswords.sync.localSyncKeypair
import com.peachpasswords.sync.rotateDeviceId
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,
    private val pendingVouchedStore: PendingVouchedDeviceStore,
    @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 credentials = loadOrCreateLocalIdentity(appContext)
        val keys = localSyncKeypair(credentials)
        val identity = buildLocalIdentity(appContext, credentials)
        VaultCrypto.zero(master)
        adoptionJob = viewModelScope.launch {
            try {
                val paired = receiver.grantFromInvitation(
                    encodedInvitation = invitation,
                    sourceDeps = sourceDeps(identity),
                    // §11.6: a v1 peer after a secure reset is reported as
                    // repair-required rather than as a generic failure.
                    onPeerOutdated = { protocolError = PEER_PROTOCOL_OUTDATED },
                    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(): AdoptionReceiverV2Deps =
        object : AdoptionReceiverV2Deps {
            override val deviceName = getDeviceName(appContext)
            override val deviceType = DeviceType.ANDROID
            override fun localIdentity(): LocalDeviceIdentity =
                buildLocalIdentity(appContext, loadOrCreateLocalIdentity(appContext))
            // Bound to the relay session by RelayAdoptionReceiver, which wraps
            // this object; the ViewModel itself holds no Noise session.
            override val pairState: AdoptionPairState? get() = null
            override suspend fun rotateLocalIdentity(): LocalDeviceIdentity {
                rotateDeviceId(appContext)
                return localIdentity()
            }
            override suspend fun validateSource(prepare: AdoptionFrame.Prepare) {}
            override suspend fun hasSourceIdentityCollision(prepare: AdoptionFrame.Prepare): Boolean =
                devices.isRevoked(prepare.deviceId, prepare.longTermPublicKeyHex)
            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 hasSourceIdentityCollision(source: AdoptionFrame.Source): Boolean =
                devices.isRevoked(source.deviceId, source.longTermPublicKeyHex)

            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 upsertPairedDeviceAndClearRevocation(device: PairedDevice) {
                devices.upsertAndClearRevocation(device)
            }

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

            override suspend fun applyLicense(licenseKey: String) {
                if (!license.storeLicense(licenseKey) || !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>, duressSlot: String?) {
                val credentials = loadOrCreateLocalIdentity(appContext)
                val adapter = VaultSyncAdapter(
                    repo = vault,
                    ourDeviceId = credentials.deviceId,
                    pairedDevices = devices,
                    localCredentials = credentials,
                )
                adapter.applyEntries(entries)
                adapter.applyTombstones(tombstones)
                if (!duressSlot.isNullOrBlank()) {
                    try {
                        val ok = vault.installAdoptedDuressSlot(duressSlot)
                        if (!ok) Log.w("PeachAdopt", "duress slot install returned false; adoption continues")
                    } catch (e: Exception) {
                        Log.w("PeachAdopt", "duress slot install failed; adoption continues", e)
                    }
                }
            }

            override suspend fun applyIncomingDeviceManifest(
                manifest: SignedDeviceManifest,
                authenticatedRemoteStaticPublicKey: ByteArray,
            ) {
                // Mesh graduation: the source's signed device graph promotes
                // the receiver into the full mesh pool. The receiver must
                // already have the source as a paired device (see applySnapshot
                // above which calls upsertPairedDevice via the Grant path),
                val credentials = loadOrCreateLocalIdentity(appContext)
                val adapter = VaultSyncAdapter(
                    repo = vault,
                    ourDeviceId = credentials.deviceId,
                    pairedDevices = devices,
                    pendingVouched = pendingVouchedStore,
                    localCredentials = credentials,
                )
                adapter.applyIncomingDeviceManifest(manifest, authenticatedRemoteStaticPublicKey)
            }

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

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

            // Bound to the relay session by RelayAdoptionHost / the receiver
            // session, which wrap this object.
            override val pairState: AdoptionPairState? get() = null

            override fun snapshotCapable(): Boolean = true

            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.storeLicense(licenseKey) || !license.isProActive()) {
                    throw IllegalStateException("license_transfer_failed")
                }
            }

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

            // The v2 holder is keyed on the validated identity, never on the
            // v1 request label that ADOPT_REQUEST_V2 does not even carry.
            override suspend fun persistPendingIdentity(device: PairedDevice, nowMs: Long): Boolean {
                val deviceName = device.deviceName ?: return false
                val deviceType = device.deviceType ?: return false
                val longTermKey = device.longTermPublicKeyHex ?: return false
                val signingKey = device.signingPublicKeyHex ?: return false
                return runCatching {
                    pendingVouchedStore.upsertAndReport(
                        PendingVouchedDevice(
                            deviceId = device.deviceId,
                            deviceName = deviceName,
                            deviceType = deviceType,
                            longTermPublicKeyHex = longTermKey,
                            signingPublicKeyHex = signingKey,
                            wsPort = device.wsPort,
                            pairedAt = nowMs,
                            receivedAt = nowMs,
                            vouchedByDeviceId = identity.deviceId,
                        ),
                    )
                }.getOrDefault(false)
            }

            override suspend fun completePendingIdentity(device: PairedDevice): Boolean =
                runCatching { pendingVouchedStore.removeIdentityAndReport(device) }.getOrDefault(false)
            override suspend fun onConsent(request: AdoptionFrame.Request): Boolean = true
            override suspend fun persistPendingHolder(request: AdoptionFrame.Request, nowMs: Long): Boolean {
                return runCatching {
                    pendingVouchedStore.upsertAndReport(
                        com.peachpasswords.sync.PendingVouchedDevice(
                            deviceId = request.sessionRequestId,
                            deviceName = request.deviceName,
                            deviceType = request.deviceType,
                            longTermPublicKeyHex = identity.longTermPublicKeyHex,
                            signingPublicKeyHex = identity.signingPublicKeyHex,
                            pairedAt = nowMs,
                            receivedAt = nowMs,
                            vouchedByDeviceId = identity.deviceId,
                        ),
                    )
                }.getOrDefault(false)
            }

            override suspend fun completePendingHolder(request: AdoptionFrame.Request, device: PairedDevice): Boolean {
                return runCatching {
                    pendingVouchedStore.removeAndReport(request.sessionRequestId)
                }.getOrDefault(false)
            }

            override fun hasDeviceIdentityCollision(device: PairedDevice): Boolean =
                devices.hasIdentityCollision(device)

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

            override suspend fun upsertPairedDevice(device: PairedDevice) {
                // SAS-confirmed direct adoption atomically validates identity,
                // lifts terminal revocation, and persists the paired row.
                devices.upsertAndClearRevocation(device)
            }

            override suspend fun vaultSnapshot(): Pair<Map<String, String>, Map<String, Long>>? {
                val credentials = loadOrCreateLocalIdentity(appContext)
                val adapter = VaultSyncAdapter(
                    repo = vault,
                    ourDeviceId = credentials.deviceId,
                    pairedDevices = devices,
                    localCredentials = credentials,
                )
                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 credentials = loadOrCreateLocalIdentity(appContext)
                val adapter = VaultSyncAdapter(
                    repo = vault,
                    ourDeviceId = credentials.deviceId,
                    pairedDevices = devices,
                    localCredentials = credentials,
                )
                return adapter.buildOutgoingDeviceManifest()
            }

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

            override suspend fun duressSlotJson(): String? =
                vault.exportDuressSlotForAdoption()
        }

    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
}