Tusker/NotificationExtension/NotificationService.swift

//
//  NotificationService.swift
//  NotificationExtension
//
//  Created by Shadowfacts on 4/9/24.
//  Copyright © 2024 Shadowfacts. All rights reserved.
//

import UserNotifications
import UserAccounts
import PushNotifications
import CryptoKit
import OSLog
import Pachyderm
import Intents
import HTMLStreamer
import UIKit
import TuskerPreferences

private let logger = Logger(subsystem: Bundle.main.bundleIdentifier!, category: "NotificationService")

private let emojiRegex = try! NSRegularExpression(pattern: ":(\\w+):", options: [])

class NotificationService: UNNotificationServiceExtension {
    
    private static let textConverter = TextConverter(configuration: .init(insertNewlines: false), callbacks: HTMLCallbacks.self)
    
    private var pendingRequest: (UNMutableNotificationContent, (UNNotificationContent) -> Void, Task<Void, Never>)?

    override func didReceive(_ request: UNNotificationRequest, withContentHandler contentHandler: @escaping (UNNotificationContent) -> Void) {
        guard let mutableContent = request.content.mutableCopy() as? UNMutableNotificationContent else {
            logger.error("Couldn't get mutable content")
            contentHandler(request.content)
            return
        }
        
        guard request.content.userInfo["v"] as? Int == 1,
              let accountID = (request.content.userInfo["ctx"] as? String).flatMap(\.removingPercentEncoding),
              let account = UserAccountsManager.shared.getAccount(id: accountID),
              let subscription = getSubscription(account: account),
              let encryptedBody = (request.content.userInfo["data"] as? String).flatMap({ Data(base64Encoded: $0) }),
              let salt = (request.content.userInfo["salt"] as? String).flatMap(decodeBase64URL(_:)),
              let serverPublicKeyData = (request.content.userInfo["pk"] as? String).flatMap(decodeBase64URL(_:)) else {
            logger.error("Missing info from push notification")
            contentHandler(request.content)
            return
        }
        
        guard let body = decryptNotification(subscription: subscription, serverPublicKeyData: serverPublicKeyData, salt: salt, encryptedBody: encryptedBody) else {
            contentHandler(request.content)
            return
        }

        let withoutPadding = body.dropFirst(2)
        
        let notification: PushNotification
        do {
            notification = try JSONDecoder().decode(PushNotification.self, from: withoutPadding)
        } catch {
            logger.error("Unable to decode push payload: \(String(describing: error))")
            contentHandler(request.content)
            return
        }
        
        mutableContent.title = notification.title
        mutableContent.body = notification.body
        mutableContent.userInfo["notificationID"] = notification.notificationID
        mutableContent.userInfo["accountID"] = accountID
        mutableContent.targetContentIdentifier = accountID

        let task = Task {
            await updateNotificationContent(mutableContent, account: account, push: notification)
            if !Task.isCancelled {
                contentHandler(pendingRequest?.0 ?? mutableContent)
                pendingRequest = nil
            }
        }
        pendingRequest = (mutableContent, contentHandler, task)
    }
    
    override func serviceExtensionTimeWillExpire() {
        if let pendingRequest {
            logger.debug("Expiring with pending request")
            pendingRequest.2.cancel()
            pendingRequest.1(pendingRequest.0)
            self.pendingRequest = nil
        } else {
            logger.debug("Expiring without pending request")
        }
    }
    
    private func updateNotificationContent(_ content: UNMutableNotificationContent, account: UserAccountInfo, push: PushNotification) async {
        let client = Client(baseURL: account.instanceURL, accessToken: account.accessToken)
        let notification: Pachyderm.Notification
        do {
            notification = try await client.run(Client.getNotification(id: push.notificationID)).0
        } catch {
            logger.error("Error fetching notification: \(String(describing: error))")
            return
        }
        
        let kindStr: String?
        switch notification.kind {
        case .reblog:
            kindStr = "🔁 Reblogged"
        case .favourite:
            kindStr = "⭐️ Favorited"
        case .follow:
            kindStr = "👤 Followed by @\(notification.account.acct)"
        case .followRequest:
            kindStr = "👤 Asked to follow by @\(notification.account.acct)"
        case .poll:
            kindStr = "📊 Poll finished"
        case .update:
            kindStr = "✏️ Edited"
        case .emojiReaction:
            if let emoji = notification.emoji {
                kindStr = "\(emoji) Reacted"
            } else {
                kindStr = nil
            }
        default:
            kindStr = nil
        }
        
        let notificationContent: String?
        if let status = notification.status {
            if notification.kind == .mention || notification.kind == .status,
               !status.spoilerText.isEmpty {
                notificationContent = "⚠️ \(status.spoilerText)"
            } else {
                notificationContent = NotificationService.textConverter.convert(html: status.content)
            }
        } else if notification.kind == .follow || notification.kind == .followRequest {
            notificationContent = nil
        } else {
            notificationContent = push.body
        }
        
        content.body = [kindStr, notificationContent].compactMap { $0 }.joined(separator: "\n")
        
        let attachmentDataTask: Task<URL?, Never>?
        // We deliberately don't include attachments for other types of notifications that have statuses (favs, etc.)
        // because we risk just fetching the same thing a bunch of times for many senders.
        if notification.kind == .mention || notification.kind == .status || notification.kind == .update,
           let status = notification.status,
           !status.sensitive,
           let attachment = status.attachments.first {
            let url = attachment.previewURL ?? attachment.url
            attachmentDataTask = Task {
                do {
                    let data = try await URLSession.shared.data(from: url).0
                    let localAttachmentURL = FileManager.default.temporaryDirectory.appendingPathComponent("attachment_\(attachment.id)").appendingPathExtension(url.pathExtension)
                    try data.write(to: localAttachmentURL)
                    return localAttachmentURL
                } catch {
                    logger.error("Error setting notification attachments: \(String(describing: error))")
                    return nil
                }
            }
        } else {
            attachmentDataTask = nil
        }
        
        let conversationIdentifier: String?
        if let status = notification.status {
            if let context = status.pleromaExtras?.context {
                conversationIdentifier = "context:\(context)"
            } else if [Notification.Kind.reblog, .favourite, .poll, .update].contains(notification.kind) {
                conversationIdentifier = "status:\(status.id)"
            } else {
                conversationIdentifier = nil
            }
        } else {
            conversationIdentifier = nil
        }
        
        let account: Account?
        switch notification.kind {
        case .mention, .status:
            account = notification.status?.account
        default:
            account = notification.account
        }
        let sender: INPerson?
        if let account {
            let handle = INPersonHandle(value: "@\(account.acct)", type: .unknown)
            let image: INImage?
            if let avatar = account.avatar,
               let (data, resp) = try? await URLSession.shared.data(from: avatar),
               let code = (resp as? HTTPURLResponse)?.statusCode,
               (200...299).contains(code) {
                image = INImage(imageData: data)
            } else {
                image = nil
            }
            sender = INPerson(
                personHandle: handle,
                nameComponents: nil,
                displayName: account.displayName,
                image: image,
                contactIdentifier: nil,
                customIdentifier: account.id
            )
        } else {
            sender = nil
        }

        let intent = INSendMessageIntent(
            recipients: nil,
            outgoingMessageType: .outgoingMessageText,
            content: notificationContent,
            speakableGroupName: nil,
            conversationIdentifier: conversationIdentifier,
            serviceName: nil,
            sender: sender,
            attachments: nil
        )
        
        let interaction = INInteraction(intent: intent, response: nil)
        interaction.direction = .incoming
        
        do {
            try await interaction.donate()
        } catch {
            logger.error("Error donating interaction: \(String(describing: error))")
            return
        }
        
        let updatedContent: UNMutableNotificationContent
       
        let contentProviding: any UNNotificationContentProviding
        if #available(iOS 18.0, visionOS 2.0, *),
           await Preferences.shared.hasFeatureFlag(.pushNotifCustomEmoji) {
            let attributedString = NSMutableAttributedString(string: content.body)
            
            for match in emojiRegex.matches(in: content.body, range: NSRange(location: 0, length: content.body.utf16.count)).reversed() {
                let emojiName = (content.body as NSString).substring(with: match.range(at: 1))
                guard let emoji = notification.status?.emojis.first(where: { $0.shortcode == emojiName }),
                      let (data, _) = try? await URLSession.shared.data(from: emoji.url),
                      let image = UIImage(data: data) else {
                    continue
                }
                let attachment = NSTextAttachment(image: image)
                let attachmentStr = NSAttributedString(attachment: attachment)
                attributedString.replaceCharacters(in: match.range, with: attachmentStr)
            }
            
            let attributedCtx = UNNotificationAttributedMessageContext(sendMessageIntent: intent, attributedContent: attributedString)
            contentProviding = attributedCtx
        } else {
            contentProviding = intent
        }
        
        do {
            let newContent = try content.updating(from: contentProviding)
            if let newMutableContent = newContent.mutableCopy() as? UNMutableNotificationContent {
                pendingRequest?.0 = newMutableContent
                updatedContent = newMutableContent
            } else {
                updatedContent = content
            }
        } catch {
            logger.error("Error updating notification from intent: \(String(describing: error))")
            updatedContent = content
        }
        
        if let localAttachmentURL = await attachmentDataTask?.value,
           let attachment = try? UNNotificationAttachment(identifier: localAttachmentURL.lastPathComponent, url: localAttachmentURL) {
            updatedContent.attachments = [
                attachment
            ]
        }
    }
    
    private func getSubscription(account: UserAccountInfo) -> PushNotifications.PushSubscription? {
        DispatchQueue.main.sync {
            // this is necessary because of a swift bug: https://github.com/apple/swift/pull/72507
            MainActor.runUnsafely {
                PushManager.shared.pushSubscription(account: account)
            }
        }
    }
    
    private func decryptNotification(subscription: PushNotifications.PushSubscription, serverPublicKeyData: Data, salt: Data, encryptedBody: Data) -> Data? {
        // See https://github.com/ClearlyClaire/webpush/blob/f14a4d52e201128b1b00245d11b6de80d6cfdcd9/lib/webpush/encryption.rb
        
        var context = Data()
        context.append(0)
        let clientPublicKey = subscription.secretKey.publicKey.x963Representation
        let clientPublicKeyLength = UInt16(clientPublicKey.count)
        context.append(UInt8((clientPublicKeyLength >> 8) & 0xFF))
        context.append(UInt8(clientPublicKeyLength & 0xFF))
        context.append(clientPublicKey)
        let serverPublicKeyLength = UInt16(serverPublicKeyData.count)
        context.append(UInt8((serverPublicKeyLength >> 8) & 0xFF))
        context.append(UInt8(serverPublicKeyLength & 0xFF))
        context.append(serverPublicKeyData)
        
        func info(encoding: String) -> Data {
            var info = Data("Content-Encoding: \(encoding)\0P-256".utf8)
            info.append(context)
            return info
        }

        let sharedSecret: SharedSecret
        do {
            let serverPublicKey = try P256.KeyAgreement.PublicKey(x963Representation: serverPublicKeyData)
            sharedSecret = try subscription.secretKey.sharedSecretFromKeyAgreement(with: serverPublicKey)
        } catch {
            logger.error("Error getting shared secret: \(String(describing: error))")
            return nil
        }
        
        let sharedInfo = Data("Content-Encoding: auth\0".utf8)
        let pseudoRandomKey = sharedSecret.hkdfDerivedSymmetricKey(using: SHA256.self, salt: subscription.authSecret, sharedInfo: sharedInfo, outputByteCount: 32)
        let contentEncryptionKeyInfo = info(encoding: "aesgcm")
        let contentEncryptionKey = HKDF<SHA256>.deriveKey(inputKeyMaterial: pseudoRandomKey, salt: salt, info: contentEncryptionKeyInfo, outputByteCount: 16)
        let nonceInfo = info(encoding: "nonce")
        let nonce = HKDF<SHA256>.deriveKey(inputKeyMaterial: pseudoRandomKey, salt: salt, info: nonceInfo, outputByteCount: 12)
        
        let nonceAndEncryptedBody = nonce.withUnsafeBytes { noncePtr in
            var data = Data(buffer: noncePtr.bindMemory(to: UInt8.self))
            data.append(encryptedBody)
            return data
        }
        do {
            let sealedBox = try AES.GCM.SealedBox(combined: nonceAndEncryptedBody)
            let decrypted = try AES.GCM.open(sealedBox, using: contentEncryptionKey)
            return decrypted
        } catch {
            logger.error("Error decrypting push: \(String(describing: error))")
            return nil
        }
    }

}

extension MainActor {
    @_unavailableFromAsync
    @available(macOS, obsoleted: 14.0)
    @available(iOS, obsoleted: 17.0)
    @available(watchOS, obsoleted: 10.0)
    @available(tvOS, obsoleted: 17.0)
    @available(visionOS 1.0, *)
    static func runUnsafely<T>(_ body: @MainActor () throws -> T) rethrows -> T {
        if #available(macOS 14.0, iOS 17.0, watchOS 10.0, tvOS 17.0, *) {
            return try MainActor.assumeIsolated(body)
        }
        
        dispatchPrecondition(condition: .onQueue(.main))
        return try withoutActuallyEscaping(body) { fn in
            try unsafeBitCast(fn, to: (() throws -> T).self)()
        }
    }
}

private func decodeBase64URL(_ s: String) -> Data? {
    var str = s.replacingOccurrences(of: "-", with: "+").replacingOccurrences(of: "_", with: "/")
    if str.count % 4 != 0 {
        str.append(String(repeating: "=", count: 4 - str.count % 4))
    }
    return Data(base64Encoded: str)
}

// copied from HTMLConverter.Callbacks, blergh
private struct HTMLCallbacks: HTMLConversionCallbacks {
    static func makeURL(string: String) -> URL? {
        try? URL.ParseStrategy().parse(string)
    }
    
    static func elementAction(name: String, attributes: [Attribute]) -> ElementAction {
        guard name == "span" else {
            return .default
        }
        let clazz = attributes.attributeValue(for: "class")
        if clazz == "invisible" {
            return .skip
        } else if clazz == "ellipsis" {
            return .append("…")
        } else {
            return .default
        }
    }
}