//
// Copyright 2025 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//

import Foundation
import LibSignalClient
import SignalServiceKit

public class QuickRestoreManager {
    public typealias RestoreMethodToken = String

    public enum Error: Swift.Error {
        case errorWaitingForNewDevice
        case invalidRegistrationMessage
        case unsupportedRestoreMethod
        case missingRestoreInformation
        case unknown
    }

    private let accountKeyStore: AccountKeyStore
    private let backupNonceStore: BackupNonceMetadataStore
    private let backupSettingsStore: BackupSettingsStore
    private let db: any DB
    private let deviceProvisioningService: DeviceProvisioningService
    private let identityManager: OWSIdentityManager
    private let networkManager: any NetworkManagerProtocol
    private let tsAccountManager: TSAccountManager

    init(
        accountKeyStore: AccountKeyStore,
        backupNonceStore: BackupNonceMetadataStore,
        backupSettingsStore: BackupSettingsStore,
        db: any DB,
        deviceProvisioningService: DeviceProvisioningService,
        identityManager: OWSIdentityManager,
        networkManager: any NetworkManagerProtocol,
        tsAccountManager: TSAccountManager,
    ) {
        self.accountKeyStore = accountKeyStore
        self.backupNonceStore = backupNonceStore
        self.backupSettingsStore = backupSettingsStore
        self.db = db
        self.deviceProvisioningService = deviceProvisioningService
        self.identityManager = identityManager
        self.networkManager = networkManager
        self.tsAccountManager = tsAccountManager
    }

    public func register(deviceProvisioningUrl: DeviceProvisioningURL) async throws -> RestoreMethodToken {
        let (
            localIdentifiers,
            accountEntropyPool,
            aciIdentityKeyPair,
            pniIdentityKeyPair,
            pinCode,
            backupTier,
            lastBackupDate,
            lastBackupSizeBytes,
            lastBackupForwardSecrecyToken,
            nextBackupSecretData,
        ) = try db.read { tx in
            guard let localIdentifiers = tsAccountManager.localIdentifiers(tx: tx) else {
                owsFailDebug("Can't quick restore without local identifiers")
                throw Error.missingRestoreInformation
            }
            guard let accountEntropyPool = accountKeyStore.getAccountEntropyPool(tx: tx) else {
                // This should be impossible; the only times you don't have
                // a AEP are during registration.
                owsFailDebug("Can't quick restore without AccountEntropyPool")
                throw Error.missingRestoreInformation
            }
            guard let aciIdentityKeyPair = identityManager.identityKeyPair(for: .aci, tx: tx) else {
                owsFailDebug("Can't quick restore without local identity key")
                throw Error.missingRestoreInformation
            }
            guard let pniIdentityKeyPair = identityManager.identityKeyPair(for: .pni, tx: tx) else {
                owsFailDebug("Can't quick restore without local identity key")
                throw Error.missingRestoreInformation
            }
            let pinCode = SSKEnvironment.shared.ows2FAManagerRef.pinCode(transaction: tx)

            let backupTier: RegistrationProvisioningMessage.BackupTier? = switch backupSettingsStore.backupPlan(tx: tx) {
            case .free: .free
            case .paid, .paidExpiringSoon, .paidAsTester: .paid
            case .disabled, .disabling: nil
            }

            let lastBackupTime: UInt64?
            let lastBackupSizeBytes: UInt64?
            if backupTier != nil {
                let lastBackupDetails = backupSettingsStore.lastBackupDetails(tx: tx)
                lastBackupTime = lastBackupDetails?.date.ows_millisecondsSince1970
                lastBackupSizeBytes = lastBackupDetails?.backupTotalSizeBytes
            } else {
                lastBackupTime = nil
                lastBackupSizeBytes = nil
            }

            let backupKey = try MessageRootBackupKey(
                accountEntropyPool: accountEntropyPool,
                aci: localIdentifiers.aci,
            )
            let lastBackupForwardSecrecyToken = try backupNonceStore.getLastForwardSecrecyToken(
                for: backupKey,
                tx: tx,
            )
            let nextBackupSecretData = backupNonceStore.getNextSecretMetadata(
                for: backupKey,
                tx: tx,
            )

            return (
                localIdentifiers,
                accountEntropyPool,
                aciIdentityKeyPair,
                pniIdentityKeyPair,
                pinCode,
                backupTier,
                lastBackupTime,
                lastBackupSizeBytes,
                lastBackupForwardSecrecyToken,
                nextBackupSecretData,
            )
        }

        let myAci = localIdentifiers.aci
        guard let myPhoneNumber = E164(localIdentifiers.phoneNumber) else {
            owsFailDebug("Can't quick restore without e164")
            throw Error.missingRestoreInformation
        }

        let restoreMethodToken = UUID().uuidString

        let registrationMessage = RegistrationProvisioningMessage(
            accountEntropyPool: accountEntropyPool,
            aci: myAci,
            aciIdentityKeyPair: aciIdentityKeyPair.identityKeyPair,
            pniIdentityKeyPair: pniIdentityKeyPair.identityKeyPair,
            phoneNumber: myPhoneNumber,
            pin: pinCode,
            tier: backupTier,
            backupVersion: BackupArchiveManagerImpl.Constants.supportedBackupVersion,
            backupTimestamp: lastBackupDate,
            backupSizeBytes: lastBackupSizeBytes,
            restoreMethodToken: restoreMethodToken,
            lastBackupForwardSecrecyToken: lastBackupForwardSecrecyToken,
            nextBackupSecretData: nextBackupSecretData,
        )

        let theirPublicKey = deviceProvisioningUrl.publicKey
        let messageBody = try registrationMessage.buildEncryptedMessageBody(theirPublicKey: theirPublicKey)
        try await deviceProvisioningService.provisionDevice(
            messageBody: messageBody,
            ephemeralDeviceId: deviceProvisioningUrl.ephemeralDeviceId,
        )

        return restoreMethodToken
    }

    public enum RestoreMethodType {
        case remoteBackup
        case localBackup
        case deviceTransfer(String)
        case decline

        fileprivate init?(response: Requests.WaitForRestoreMethodChoice.Response) {
            switch response.method {
            case .decline: self = .decline
            case .localBackup: self = .localBackup
            case .remoteBackup: self = .remoteBackup
            case .deviceTransfer:
                guard let bootstrapData = response.deviceTransferBootstrap else { return nil }
                self = .deviceTransfer(bootstrapData)
            }
        }
    }

    public func reportRestoreMethodChoice(method: RestoreMethodType, restoreMethodToken: RestoreMethodToken) async throws {
        whileLoop: while true {
            let response = try await networkManager.asyncRequest(
                Requests.ChooseRestoreMethod.buildRequest(
                    token: restoreMethodToken,
                    method: method,
                ),
            )
            switch response.responseStatusCode {
            case 200, 204:
                return
            case 429:
                try await Task.sleep(
                    nanoseconds: HTTPUtils.retryDelayNanoSeconds(response, defaultRetryTime: Constants.defaultRetryTime),
                )
                continue whileLoop
            default:
                owsFailDebug("Unexpected response")
                throw Error.unknown
            }
        }
    }

    public func waitForRestoreMethodChoice(restoreMethodToken: RestoreMethodToken) async throws -> RestoreMethodType {
        whileLoop: while true {
            do {
                let response = try await networkManager.asyncRequest(
                    Requests.WaitForRestoreMethodChoice.buildRequest(token: restoreMethodToken),
                )
                switch response.responseStatusCode {
                case 200:
                    guard
                        let data = response.responseBodyData,
                        let response = try? JSONDecoder().decode(
                            Requests.WaitForRestoreMethodChoice.Response.self,
                            from: data,
                        )
                    else {
                        throw Error.errorWaitingForNewDevice
                    }

                    guard let responseType = RestoreMethodType(response: response) else {
                        throw Error.unsupportedRestoreMethod
                    }
                    return responseType
                case 400:
                    throw Error.invalidRegistrationMessage
                case 204:
                    /// The timeout elapsed without the device linking; clients can request again.
                    continue whileLoop
                case 429:
                    try await Task.sleep(
                        nanoseconds: HTTPUtils.retryDelayNanoSeconds(response, defaultRetryTime: Constants.defaultRetryTime),
                    )
                    continue whileLoop
                default:
                    owsFailDebug("Unexpected response")
                    throw Error.unknown
                }
            } catch {
                owsFailDebug("Unexpected exception")
                throw Error.unknown
            }
        }
    }

    private enum Constants {
        static let longPollRequestTimeoutSeconds: UInt32 = 60 * 5
        static let defaultRetryTime: TimeInterval = 15
    }

    fileprivate enum Requests {
        enum RestoreMethod: String, Codable {
            case remoteBackup = "REMOTE_BACKUP"
            case localBackup = "LOCAL_BACKUP"
            case deviceTransfer = "DEVICE_TRANSFER"
            case decline = "DECLINE"
        }

        enum WaitForRestoreMethodChoice {
            struct Response: Codable {
                /// The method of restore chosen by the new device
                let method: RestoreMethod
                /// Additional data used to bootstrap device transfer
                let deviceTransferBootstrap: String?
            }

            static func buildRequest(token: RestoreMethodToken) -> TSRequest {
                var urlComponents = URLComponents(string: "v1/devices/restore_account/\(token)")!
                urlComponents.queryItems = [URLQueryItem(
                    name: "timeout",
                    value: "\(Constants.longPollRequestTimeoutSeconds)",
                )]
                var request = TSRequest(
                    url: urlComponents.url!,
                    method: "GET",
                    parameters: nil,
                )

                request.auth = .anonymous
                request.applyRedactionStrategy(.redactURL())
                // The timeout is server side; apply wiggle room for our local clock.
                request.timeoutInterval = 10 + TimeInterval(Constants.longPollRequestTimeoutSeconds)
                return request
            }
        }

        enum ChooseRestoreMethod {
            static func buildRequest(token: RestoreMethodToken, method: RestoreMethodType) -> TSRequest {
                var deviceTransferBootstrap: String?
                let method: RestoreMethod = {
                    switch method {
                    case .decline: return .decline
                    case .deviceTransfer(let data):
                        deviceTransferBootstrap = data
                        return .deviceTransfer
                    case .remoteBackup:
                        return .remoteBackup
                    case .localBackup:
                        return .localBackup
                    }
                }()

                var parameters: [String: Any] = ["method": method.rawValue]
                // `deviceTransferBootstrap` contains unpadded base64 encoded data that is used by
                // the other device to initiate device transfer. Note that server enforces a
                // 4096 bytes limit on this field.
                deviceTransferBootstrap.map { parameters["deviceTransferBootstrap"] = $0 }

                let urlComponents = URLComponents(string: "v1/devices/restore_account/\(token)")!
                var request = TSRequest(
                    url: urlComponents.url!,
                    method: "PUT",
                    parameters: parameters,
                )

                request.auth = .anonymous
                request.applyRedactionStrategy(.redactURL())
                return request
            }
        }
    }
}
