Tusker/Packages/PushNotifications/Sources/PushNotifications/PushManagerImpl.swift

//
//  PushManagerImpl.swift
//  PushNotifications
//
//  Created by Shadowfacts on 4/7/24.
//

import UIKit
import UserAccounts
import CryptoKit

class PushManagerImpl: _PushManager {
    private let endpoint: URL
    
    var enabled: Bool {
        true
    }
    
    private var apnsEnvironment: String {
        #if DEBUG
        "development"
        #else
        "production"
        #endif
    }
    
    private var remoteNotificationsRegistrationContinuation: CheckedContinuation<Data, any Error>?
    
    private let defaults = UserDefaults(suiteName: "group.space.vaccor.Tusker")!
    public private(set) var subscriptions: [PushSubscription] {
        get {
            if let array = defaults.array(forKey: "PushSubscriptions") as? [[String: Any]] {
                return array.compactMap(PushSubscription.init(defaultsDict:))
            } else {
                return []
            }
        }
        set {
            defaults.setValue(newValue.map(\.defaultsDict), forKey: "PushSubscriptions")
        }
    }
    
    init(endpoint: URL) {
        self.endpoint = endpoint
    }
    
    func createSubscription(account: UserAccountInfo) async throws -> PushSubscription {
        if let existing = pushSubscription(account: account) {
            return existing
        }
        let key = P256.KeyAgreement.PrivateKey()
        var authSecret = Data(count: 16)
        let res = authSecret.withUnsafeMutableBytes { ptr in
            SecRandomCopyBytes(kSecRandomDefault, 16, ptr.baseAddress!)
        }
        guard res == errSecSuccess else {
            throw CreateSubscriptionError.generatingAuthSecret(res)
        }
        let token = try await getDeviceToken()
        let subscription = PushSubscription(
            accountID: account.id,
            endpoint: endpointURL(deviceToken: token, accountID: account.id),
            secretKey: key,
            authSecret: authSecret,
            alerts: [],
            policy: .all
        )
        subscriptions.append(subscription)
        return subscription
    }
    
    private func endpointURL(deviceToken: Data, accountID: String) -> URL {
        var endpoint = URLComponents(url: endpoint, resolvingAgainstBaseURL: false)!
        let accountID = accountID.addingPercentEncoding(withAllowedCharacters: .alphanumerics)!
        endpoint.path = "/push/v1/\(apnsEnvironment)/\(deviceToken.hexEncodedString())/\(accountID)"
        return endpoint.url!
    }
    
    func removeSubscription(account: UserAccountInfo) {
        subscriptions.removeAll { $0.accountID == account.id }
    }
    
    func updateSubscription(account: UserAccountInfo, alerts: PushSubscription.Alerts, policy: PushSubscription.Policy) {
        guard let index = subscriptions.firstIndex(where: { $0.accountID == account.id }) else {
            return
        }
        var copy = subscriptions[index]
        copy.alerts = alerts
        copy.policy = policy
        subscriptions[index] = copy
    }
    
    func pushSubscription(account: UserAccountInfo) -> PushSubscription? {
        subscriptions.first { $0.accountID == account.id }
    }
    
    func updateIfNecessary(updateSubscription: @escaping (PushSubscription) async -> Bool) async {
        let subscriptions = self.subscriptions
        guard !subscriptions.isEmpty else {
            return
        }
        do {
            let token = try await getDeviceToken()
            self.subscriptions = await AsyncSequenceAdaptor(wrapping: subscriptions).map {
                let newEndpoint = await self.endpointURL(deviceToken: token, accountID: $0.accountID)
                guard newEndpoint != $0.endpoint else {
                    PushManager.logger.debug("Skipping update of push subscription with endpoint \($0.endpoint, privacy: .public)")
                    return $0
                }
                var copy = $0
                copy.endpoint = newEndpoint
                if await updateSubscription(copy) {
                    return copy
                } else {
                    return $0
                }
            }.reduce(into: [], { partialResult, el in
                partialResult.append(el)
            })
        } catch {
            PushManager.logger.error("Failed to update push registration: \(String(describing: error), privacy: .public)")
            PushManager.captureError?(error)
        }
    }
    
    private func getDeviceToken() async throws -> Data {
        return try await withCheckedThrowingContinuation { continuation in
            remoteNotificationsRegistrationContinuation = continuation
            UIApplication.shared.registerForRemoteNotifications()
        }
    }
    
    func didRegisterForRemoteNotifications(deviceToken: Data) {
        remoteNotificationsRegistrationContinuation?.resume(returning: deviceToken)
        remoteNotificationsRegistrationContinuation = nil
    }
    
    func didFailToRegisterForRemoteNotifications(error: any Error) {
        remoteNotificationsRegistrationContinuation?.resume(throwing: PushRegistrationError.registeringForRemoteNotifications(error))
        remoteNotificationsRegistrationContinuation = nil
    }
}

enum PushRegistrationError: LocalizedError {
    case alreadyRegistering
    case registeringForRemoteNotifications(any Error)
    
    var errorDescription: String? {
        switch self {
        case .alreadyRegistering:
            "Already registering"
        case .registeringForRemoteNotifications(let error):
            "Remote notifications: \(error.localizedDescription)"
        }
    }
}

enum CreateSubscriptionError: LocalizedError {
    case generatingAuthSecret(OSStatus)
    
    var errorDescription: String? {
        switch self {
        case .generatingAuthSecret(let code):
            "Generating auth secret: \(code)"
        }
    }
}

private extension Data {
    func hexEncodedString() -> String {
        String(unsafeUninitializedCapacity: count * 2) { buffer in
            let chars = Array("0123456789ABCDEF".utf8)
            for (i, x) in enumerated() {
                let (upper, lower) = x.quotientAndRemainder(dividingBy: 16)
                buffer[i * 2] = chars[Int(upper)]
                buffer[i * 2 + 1] = chars[Int(lower)]
            }
            return count * 2
        }
    }
}

private struct AsyncSequenceAdaptor<S: Sequence>: AsyncSequence {
    typealias Element = S.Element
    
    let base: S
    
    init(wrapping base: S) {
        self.base = base
    }
    
    func makeAsyncIterator() -> AsyncIterator {
        AsyncIterator(base: base.makeIterator())
    }
    
    struct AsyncIterator: AsyncIteratorProtocol {
        var base: S.Iterator
        
        mutating func next() async -> Element? {
            base.next()
        }
    }
}