/*
 * Copyright 2023, 2024 New Vector Ltd.
 *
 * SPDX-License-Identifier: AGPL-3.0-only
 * Please see LICENSE in the repository root for full details.
 */

package io.element.android.libraries.matrix.impl.verification

import io.element.android.libraries.core.data.tryOrNull
import io.element.android.libraries.matrix.api.verification.SessionVerificationData
import io.element.android.libraries.matrix.api.verification.SessionVerificationRequestDetails
import io.element.android.libraries.matrix.api.verification.SessionVerificationService
import io.element.android.libraries.matrix.api.verification.SessionVerificationServiceListener
import io.element.android.libraries.matrix.api.verification.SessionVerifiedStatus
import io.element.android.libraries.matrix.api.verification.VerificationEmoji
import io.element.android.libraries.matrix.api.verification.VerificationFlowState
import io.element.android.libraries.matrix.impl.util.cancelAndDestroy
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.NonCancellable
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.SharingStarted
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.flow.stateIn
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeout
import org.matrix.rustcomponents.sdk.Client
import org.matrix.rustcomponents.sdk.Encryption
import org.matrix.rustcomponents.sdk.RecoveryState
import org.matrix.rustcomponents.sdk.RecoveryStateListener
import org.matrix.rustcomponents.sdk.SessionVerificationController
import org.matrix.rustcomponents.sdk.SessionVerificationControllerDelegate
import org.matrix.rustcomponents.sdk.VerificationState
import org.matrix.rustcomponents.sdk.VerificationStateListener
import org.matrix.rustcomponents.sdk.use
import timber.log.Timber
import kotlin.time.Duration.Companion.seconds
import org.matrix.rustcomponents.sdk.SessionVerificationData as RustSessionVerificationData
import org.matrix.rustcomponents.sdk.SessionVerificationRequestDetails as RustSessionVerificationRequestDetails

class RustSessionVerificationService(
    private val client: Client,
    isSyncServiceReady: Flow<Boolean>,
    private val sessionCoroutineScope: CoroutineScope,
) : SessionVerificationService, SessionVerificationControllerDelegate {
    private val encryptionService: Encryption = client.encryption()
    private lateinit var verificationController: SessionVerificationController

    private val _verificationFlowState = MutableStateFlow<VerificationFlowState>(VerificationFlowState.Initial)
    override val verificationFlowState = _verificationFlowState.asStateFlow()

    private val _sessionVerifiedStatus = MutableStateFlow<SessionVerifiedStatus>(SessionVerifiedStatus.Unknown)
    override val sessionVerifiedStatus: StateFlow<SessionVerifiedStatus> = _sessionVerifiedStatus.asStateFlow()

    private val recoveryState = MutableStateFlow(RecoveryState.UNKNOWN)

    // Listen for changes in verification status and update accordingly
    private val verificationStateListenerTaskHandle = encryptionService.verificationStateListener(object : VerificationStateListener {
        override fun onUpdate(status: VerificationState) {
            Timber.d("New verification state: $status")
            _sessionVerifiedStatus.value = status.map()
        }
    })

    // In case we enter the recovery key instead we check changes in the recovery state, since the listener above won't be triggered
    private val recoveryStateListenerTaskHandle = encryptionService.recoveryStateListener(object : RecoveryStateListener {
        override fun onUpdate(status: RecoveryState) {
            Timber.d("New recovery state: $status")
            // We could check the `RecoveryState`, but it's easier to just use the verification state directly
            recoveryState.value = status
        }
    })

    /**
     * The internal service that checks verification can only run after the initial sync.
     * This [StateFlow] will notify consumers when the service is ready to be used.
     */
    private val isReady = isSyncServiceReady.stateIn(sessionCoroutineScope, SharingStarted.Eagerly, false)

    override val needsSessionVerification = sessionVerifiedStatus.map { verificationStatus ->
        verificationStatus == SessionVerifiedStatus.NotVerified
    }

    private var isOwnVerification = true

    override fun didReceiveVerificationRequest(details: RustSessionVerificationRequestDetails) {
        listener?.onIncomingSessionRequest(details.map())
    }

    private var listener: SessionVerificationServiceListener? = null

    init {
        // Instantiate the verification controller when possible, this is needed to get incoming verification requests
        sessionCoroutineScope.launch {
            // Needed to avoid crashes on unit tests due to the Rust SDK not being available
            tryOrNull {
                encryptionService.waitForE2eeInitializationTasks()
                initVerificationControllerIfNeeded()
            }
        }
    }

    override fun setListener(listener: SessionVerificationServiceListener?) {
        this.listener = listener
    }

    override suspend fun requestVerification() = tryOrFail {
        initVerificationControllerIfNeeded()
        verificationController.requestVerification()
    }

    override suspend fun cancelVerification() = tryOrFail {
        verificationController.cancelVerification()
        // We need to manually set the state to canceled, as the Rust SDK doesn't always call `didCancel` when it should
        didCancel()
    }

    override suspend fun approveVerification() = tryOrFail { verificationController.approveVerification() }

    override suspend fun declineVerification() = tryOrFail { verificationController.declineVerification() }

    override suspend fun startVerification() = tryOrFail {
        verificationController.startSasVerification()
    }

    override suspend fun acknowledgeVerificationRequest(details: SessionVerificationRequestDetails) = tryOrFail {
        isOwnVerification = false
        initVerificationControllerIfNeeded()
        verificationController.acknowledgeVerificationRequest(
            senderId = details.senderId.value,
            flowId = details.flowId.value,
        )
    }

    override suspend fun acceptVerificationRequest() = tryOrFail {
        verificationController.acceptVerificationRequest()
    }

    private suspend fun tryOrFail(block: suspend () -> Unit) {
        runCatching {
            // Ensure the block cannot be cancelled, else if the Rust SDK emit a new state during the API execution,
            // the state machine may cancel the api call.
            withContext(NonCancellable) {
                block()
            }
        }.onFailure {
            Timber.e(it, "Failed to verify session")
            didFail()
        }
    }

    // region Delegate implementation

    // When verification attempt is accepted by the other device
    override fun didAcceptVerificationRequest() {
        _verificationFlowState.value = VerificationFlowState.DidAcceptVerificationRequest
    }

    override fun didCancel() {
        _verificationFlowState.value = VerificationFlowState.DidCancel
    }

    override fun didFail() {
        Timber.e("Session verification failed with an unknown error")
        _verificationFlowState.value = VerificationFlowState.DidFail
    }

    override fun didFinish() {
        sessionCoroutineScope.launch {
            // Ideally this should be `verificationController?.isVerified().orFalse()` but for some reason it returns false if run immediately
            // It also sometimes unexpectedly fails to report the session as verified, so we have to handle that possibility and fail if needed
            runCatching {
                withTimeout(20.seconds) {
                    // Wait until the SDK reports the state as verified
                    sessionVerifiedStatus.first { it == SessionVerifiedStatus.Verified }
                }
            }
                .onSuccess {
                    if (isOwnVerification) {
                        // Try waiting for the final recovery state for better UX, but don't block the verification state on it
                        tryOrNull {
                            withTimeout(10.seconds) {
                                // Wait until the recovery state is either fully loaded or we check it's explicitly disabled
                                recoveryState.first { it == RecoveryState.ENABLED || it == RecoveryState.DISABLED }
                            }
                        }
                    }

                    _verificationFlowState.value = VerificationFlowState.DidFinish
                    updateVerificationStatus()
                }
                .onFailure {
                    Timber.e(it, "Verification finished, but the Rust SDK still reports the session as unverified.")
                    didFail()
                }
        }
    }

    override fun didReceiveVerificationData(data: RustSessionVerificationData) {
        _verificationFlowState.value = VerificationFlowState.DidReceiveVerificationData(data.map())
    }

    // When the actual SAS verification starts
    override fun didStartSasVerification() {
        _verificationFlowState.value = VerificationFlowState.DidStartSasVerification
    }

    // end-region

    override suspend fun reset(cancelAnyPendingVerificationAttempt: Boolean) {
        isOwnVerification = true
        if (isReady.value && cancelAnyPendingVerificationAttempt) {
            // Cancel any pending verification attempt
            tryOrNull { verificationController.cancelVerification() }
        }
        _verificationFlowState.value = VerificationFlowState.Initial
    }

    fun destroy() {
        Timber.d("Destroying RustSessionVerificationService")
        verificationStateListenerTaskHandle.cancelAndDestroy()
        recoveryStateListenerTaskHandle.cancelAndDestroy()
        if (this::verificationController.isInitialized) {
            verificationController.setDelegate(null)
        }
    }

    private var initControllerMutex = Mutex()

    private suspend fun initVerificationControllerIfNeeded() = initControllerMutex.withLock {
        if (!this::verificationController.isInitialized) {
            tryOrFail {
                verificationController = client.getSessionVerificationController()
                verificationController.setDelegate(this)
            }
        }
    }

    private fun updateVerificationStatus() {
        runCatching {
            _sessionVerifiedStatus.value = encryptionService.verificationState().map()
            Timber.d("New verification status: ${_sessionVerifiedStatus.value}")
        }
    }
}

private fun VerificationState.map() = when (this) {
    VerificationState.UNKNOWN -> SessionVerifiedStatus.Unknown
    VerificationState.VERIFIED -> SessionVerifiedStatus.Verified
    VerificationState.UNVERIFIED -> SessionVerifiedStatus.NotVerified
}

private fun RustSessionVerificationData.map(): SessionVerificationData {
    return use { sessionVerificationData ->
        when (sessionVerificationData) {
            is RustSessionVerificationData.Emojis -> {
                SessionVerificationData.Emojis(
                    emojis = sessionVerificationData.emojis.mapIndexed { index, emoji ->
                        emoji.use { sessionVerificationEmoji ->
                            VerificationEmoji(
                                number = sessionVerificationData.indices[index].toInt(),
                                emoji = sessionVerificationEmoji.symbol(),
                                description = sessionVerificationEmoji.description(),
                            )
                        }
                    },
                )
            }
            is RustSessionVerificationData.Decimals -> {
                SessionVerificationData.Decimals(
                    decimals = sessionVerificationData.values.map { it.toInt() },
                )
            }
        }
    }
}
