element-x-ios/ElementX/Sources/Screens/QRCodeLoginScreen/QRCodeLoginScreenViewModel.swift

//
// Copyright 2025 Element Creations Ltd.
// Copyright 2022-2025 New Vector Ltd.
//
// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
// Please see LICENSE files in the repository root for full details.
//

import Combine
import Foundation

typealias QRCodeLoginScreenViewModelType = StateStoreViewModel<QRCodeLoginScreenViewState, QRCodeLoginScreenViewAction>

class QRCodeLoginScreenViewModel: QRCodeLoginScreenViewModelType, QRCodeLoginScreenViewModelProtocol {
    private let appMediator: AppMediatorProtocol
    
    private let actionsSubject: PassthroughSubject<QRCodeLoginScreenViewModelAction, Never> = .init()
    var actionsPublisher: AnyPublisher<QRCodeLoginScreenViewModelAction, Never> {
        actionsSubject.eraseToAnyPublisher()
    }
    
    private var currentTask: AnyCancellable?
    private var oAuthResultTask: AnyCancellable?
    
    init(mode: QRCodeLoginScreenMode,
         canSignInManually: Bool,
         appMediator: AppMediatorProtocol) {
        self.appMediator = appMediator
        
        let initialViewState: QRCodeLoginState = switch mode {
        case .login: .loginInstructions
        case .linkDesktop: .linkDesktopInstructions
        case .linkMobile(let progressPublisher, _):
            switch progressPublisher.value {
            case .qrReady(let image): .displayQR(.active(image))
            default: .error(.unknown)
            }
        }
        
        super.init(initialViewState: .init(state: initialViewState, mode: mode, canSignInManually: canSignInManually))
        setupSubscriptions()
        
        if case .linkMobile(let progressPublisher, _) = mode {
            listenToDisplayQRProgress(progressPublisher: progressPublisher)
        }
    }
    
    // MARK: - Public
    
    override func process(viewAction: QRCodeLoginScreenViewAction) {
        switch viewAction {
        case .cancel, .errorAction(.cancel):
            actionsSubject.send(.cancel)
        case .startScan:
            Task { await startScanIfPossible() }
        case .sendCheckCode:
            Task { await sendCheckCode() }
        case .errorAction(.startOver):
            switch state.mode {
            case .login:
                // Login restarts the flow on this screen.
                state.bindings.qrResult = nil
                state.state = .loginInstructions
            case .linkDesktop, .linkMobile:
                // Linking a new device starts on the previous screen.
                actionsSubject.send(.startOver)
            }
        case .errorAction(.openSettings):
            appMediator.openAppSettings()
        case .errorAction(.signInManually):
            actionsSubject.send(.signInManually)
        }
    }
    
    // MARK: - Private
    
    private func setupSubscriptions() {
        context.$viewState
            // not using compactMap before remove duplicates because if there is an error, and the same
            // code needs to be rescanned the transition to nil to clean the state would get ignored.
            .map(\.bindings.qrResult)
            .removeDuplicates()
            .compactMap { $0 }
            // this needs to be received on the main actor or the state change for connecting won't work properly
            .receive(on: DispatchQueue.main)
            .sink { [weak self] qrData in
                guard let self else { return }
                switch state.mode {
                case .login(let qrCodeLoginService):
                    handleScan(qrData: qrData, loginService: qrCodeLoginService)
                case .linkDesktop(let linkNewDeviceService):
                    handleScan(qrData: qrData, linkService: linkNewDeviceService)
                case .linkMobile:
                    fatalError("A code should never be scanned when showing one.")
                }
            }
            .store(in: &cancellables)
    }
    
    private func startScanIfPossible() async {
        state.bindings.qrResult = nil
        state.state = await appMediator.requestAuthorizationIfNeeded() ? .scan(.scanning) : .error(.noCameraPermission)
    }
    
    private func handleScan(qrData: Data, loginService: QRCodeLoginServiceProtocol) {
        guard currentTask == nil else { return }
        
        state.state = .scan(.connecting)
        
        MXLog.info("Login scanning QR code")
        let progressPublisher = loginService.loginWithQRCode(data: qrData)
        
        currentTask = progressPublisher
            .removeDuplicates()
            .receive(on: DispatchQueue.main)
            .sink { [weak self] completion in
                guard let self else { return }
                currentTask = nil
                
                switch completion {
                case .finished: break
                case .failure(.qrCodeError(let error)):
                    handleError(error)
                case .failure:
                    handleError(.unknown)
                }
            } receiveValue: { [weak self] progress in
                MXLog.info("QR Login Progress changed to: \(progress)")
                
                guard let self,
                      // Let's not advance the state if the current state is already invalid
                      !state.state.isError else {
                    return
                }
                
                switch progress {
                case .starting:
                    break // Nothing to do, the state was set above.
                case .establishingSecureChannel(_, let stringCode):
                    state.state = .displayCode(.deviceCode(stringCode))
                case .waitingForToken(let code):
                    state.state = .displayCode(.verificationCode(code))
                case .syncingSecrets:
                    break // Nothing to do.
                case .signedIn(let session):
                    MXLog.info("QR Login completed")
                    actionsSubject.send(.signedIn(userSession: session))
                }
            }
    }
    
    private func handleScan(qrData: Data, linkService: LinkNewDeviceServiceProtocol) {
        guard currentTask == nil else { return }
        
        state.state = .scan(.connecting)
        
        MXLog.info("Link scanning QR code")
        let progressPublisher = linkService.linkDesktopDevice(with: qrData)
        
        currentTask = progressPublisher
            .removeDuplicates()
            .receive(on: DispatchQueue.main)
            .sink { [weak self] completion in
                guard let self else { return }
                currentTask = nil
                
                if case .failure(let error) = completion {
                    handleError(error)
                }
            } receiveValue: { [weak self] progress in
                MXLog.info("Linking with QR progress changed to: \(progress)")
                
                guard let self,
                      // Let's not advance the state if the current state is already invalid
                      !state.state.isError else {
                    return
                }
                
                switch progress {
                case .starting:
                    break // Nothing to do, the state was set above.
                case .establishingSecureChannel(let checkCodeString):
                    state.state = .displayCode(.deviceCode(checkCodeString))
                case .waitingForAuthorisation(let url):
                    requestOAuthAuthorization(url: url)
                case .syncingSecrets:
                    break // Nothing to do.
                case .done:
                    MXLog.info("Link with QR code completed.")
                    actionsSubject.send(.linkedDevice)
                }
            }
    }
    
    private func listenToDisplayQRProgress(progressPublisher: LinkNewDeviceService.LinkMobileProgressPublisher) {
        state.bindings.qrResult = nil
            
        MXLog.info("Link showing QR code.")
        
        currentTask = progressPublisher
            .removeDuplicates()
            .receive(on: DispatchQueue.main)
            .sink { [weak self] completion in
                guard let self else { return }
                currentTask = nil
                
                if case .failure(let error) = completion {
                    handleError(error)
                }
            } receiveValue: { [weak self] progress in
                MXLog.info("Linking with QR progress changed to: \(progress)")
                
                guard let self,
                      // Let's not advance the state if the current state is already invalid
                      !state.state.isError else {
                    return
                }
                
                switch progress {
                case .starting, .qrReady:
                    break // Nothing to do, we are already showing the code by the time this method is called.
                case .qrScanned(let checkCodeSender):
                    state.state = .confirmCode(.inputCode(checkCodeSender))
                case .waitingForAuthorisation(let url):
                    requestOAuthAuthorization(url: url)
                case .syncingSecrets:
                    break // Nothing to do.
                case .done:
                    MXLog.info("Link with QR code completed.")
                    actionsSubject.send(.linkedDevice)
                }
            }
    }
    
    private func regenerateQRCode() async {
        guard case .linkMobile(_, let clientProxy) = state.mode else {
            fatalError("Cannot display a QR code in this mode.")
        }
        
        let linkNewDeviceService = clientProxy.linkNewDeviceService()
        let progressPublisher = linkNewDeviceService.linkMobileDevice()
        
        do {
            async let minimumDelay = Task.sleep(for: .seconds(1)) // To show the spinner for long enough so there isn't a flicker.
            async let qrReadyProgress = progressPublisher.values
                .first { progress in
                    switch progress {
                    case .qrReady: true
                    default: false
                    }
                }

            guard case .qrReady(let image) = try await (qrReadyProgress, minimumDelay).0 else {
                state.state = .error(.unknown)
                return
            }

            state.state = .displayQR(.active(image))
        } catch let error as QRCodeLoginError {
            handleError(error)
        } catch {
            handleError(.unknown)
        }
        
        listenToDisplayQRProgress(progressPublisher: progressPublisher)
    }
    
    private func sendCheckCode() async {
        guard case let .confirmCode(.inputCode(checkCodeSender)) = state.state else {
            fatalError("Attempting to check code from the wrong state.")
        }
        
        state.state = .confirmCode(.sendingCode)
        
        let stringValue = state.bindings.checkCodeInput
        let code = UInt8(stringValue) ?? 0
        
        do {
            MXLog.info("Valid code entered, sending.")
            try await checkCodeSender.send(code: code)
        } catch {
            MXLog.error("Failed to send check code: \(error)")
            handleError(.unknown)
        }
    }
    
    private func requestOAuthAuthorization(url: URL) {
        let (stream, continuation) = AsyncStream<Result<Void, OAuthError>>.makeStream()
        actionsSubject.send(.requestOAuthAuthorisation(url, continuation))
        
        oAuthResultTask = Task { [weak self] in
            for await result in stream {
                guard let self else { return }
                switch result {
                case .success:
                    break // The state will be updated by the status publisher.
                case .failure(.userCancellation):
                    MXLog.info("User cancelled the WAS, the flow must be restarted.")
                    actionsSubject.send(.startOver)
                case .failure:
                    handleError(.unknown)
                }
            }
        }
        .asCancellable()
    }
    
    private func handleError(_ error: QRCodeLoginError) {
        MXLog.error("Failed to scan the QR code: \(error)")
        switch error {
        case .invalidQRCode:
            state.state = .scan(.scanFailed(.invalid))
        case .providerNotAllowed(let scannedProvider, let allowedProviders):
            state.state = .scan(.scanFailed(.notAllowed(scannedProvider: scannedProvider, allowedProviders: allowedProviders)))
        case .deviceNotSignedIn:
            state.state = .scan(.scanFailed(.deviceNotSignedIn))
        case .cancelled:
            state.state = .error(.cancelled)
        case .connectionInsecure:
            state.state = .error(.connectionNotSecure)
        case .invalidCheckCode:
            state.state = .confirmCode(.invalidCode)
        case .declined:
            state.state = .error(.declined)
        case .linkingNotSupported:
            state.state = .error(.linkingNotSupported)
        case .expired where state.state.isDisplayQR:
            state.state = .displayQR(.expired)
            Task { await regenerateQRCode() }
        case .expired:
            state.state = .error(.expired)
        case .slidingSyncNotAvailable:
            state.state = .error(.slidingSyncNotAvailable)
        case .deviceAlreadySignedIn:
            state.state = .error(.deviceAlreadySignedIn)
        case .unknown:
            state.state = .error(.unknown)
        }
    }
        
    /// Only for mocking initial states
    fileprivate init(state: QRCodeLoginState, mode: QRCodeLoginScreenMode, canSignInManually: Bool, checkCodeInput: String) {
        appMediator = AppMediatorMock(.init())
        super.init(initialViewState: .init(state: state,
                                           mode: mode,
                                           canSignInManually: canSignInManually,
                                           bindings: .init(checkCodeInput: checkCodeInput)))
    }
}

extension QRCodeLoginScreenViewModel {
    static func mock(state: QRCodeLoginState,
                     mode: QRCodeLoginScreenMode = .login(QRCodeLoginServiceMock()),
                     canSignInManually: Bool = true,
                     checkCodeInput: String = "") -> QRCodeLoginScreenViewModel {
        QRCodeLoginScreenViewModel(state: state,
                                   mode: mode,
                                   canSignInManually: canSignInManually,
                                   checkCodeInput: checkCodeInput)
    }
}