// Copyright © 2020 Metabolist. All rights reserved.

import CryptoKit
import Keychain
import Kingfisher
import Mastodon
import Secrets
import ServiceLayer
import UserNotifications

final class NotificationService: UNNotificationServiceExtension {
    private let environment = AppEnvironment.live(
        userNotificationCenter: .current(),
        reduceMotion: { false })

    override init() {
        super.init()
    }

    var contentHandler: ((UNNotificationContent) -> Void)?
    var bestAttemptContent: UNMutableNotificationContent?

    override func didReceive(
        _ request: UNNotificationRequest,
        withContentHandler contentHandler: @escaping (UNNotificationContent) -> Void) {
        self.contentHandler = contentHandler
        bestAttemptContent = (request.content.mutableCopy() as? UNMutableNotificationContent)

        guard let bestAttemptContent = bestAttemptContent else { return }

        let pushNotification: PushNotification
        let decryptedJSON: Data
        let identityId: Identity.Id

        do {
            (decryptedJSON, identityId) = try Self.extractAndDecrypt(userInfo: request.content.userInfo)

            pushNotification = try MastodonDecoder().decode(PushNotification.self, from: decryptedJSON)
        } catch {
            contentHandler(bestAttemptContent)

            return
        }

        bestAttemptContent.title = pushNotification.title
        bestAttemptContent.body = XMLUnescaper(string: pushNotification.body).unescape()

        let appPreferences = AppPreferences(environment: environment)

        if appPreferences.notificationSounds.contains(pushNotification.notificationType) {
            bestAttemptContent.sound = .default
        }

        if appPreferences.notificationAccountName,
           let accountName = try? AllIdentitiesService(environment: environment).identity(id: identityId)?.handle {
            bestAttemptContent.subtitle = accountName
        }

        if appPreferences.notificationPictures {
            Self.addImage(pushNotification: pushNotification,
                          bestAttemptContent: bestAttemptContent,
                          contentHandler: contentHandler)
        } else {
            contentHandler(bestAttemptContent)
        }
    }

    override func serviceExtensionTimeWillExpire() {
        if let contentHandler = contentHandler, let bestAttemptContent =  bestAttemptContent {
            contentHandler(bestAttemptContent)
        }
    }
}

enum NotificationServiceError: Error {
    case userInfoDataAbsent
    case keychainDataAbsent
}

private extension NotificationService {
    static let identityIdUserInfoKey = "i"
    static let encryptedMessageUserInfoKey = "m"
    static let saltUserInfoKey = "s"
    static let serverPublicKeyUserInfoKey = "k"
    static let keyLength = 16
    static let nonceLength = 12
    static let pseudoRandomKeyLength = 32
    static let paddedByteCount = 2
    static let curve = "P-256"

    enum HKDFInfo: String {
        case auth, aesgcm, nonce

        var bytes: [UInt8] {
            Array("Content-Encoding: \(self)\0".utf8)
        }
    }

    static func extractAndDecrypt(userInfo: [AnyHashable: Any]) throws -> (Data, Identity.Id) {
        guard
            let identityIdString = userInfo[identityIdUserInfoKey] as? String,
            let identityId = Identity.Id(uuidString: identityIdString),
            let encryptedMessageBase64 = (userInfo[encryptedMessageUserInfoKey] as? String)?.URLSafeBase64ToBase64(),
            let encryptedMessage = Data(base64Encoded: encryptedMessageBase64),
            let saltBase64 = (userInfo[saltUserInfoKey] as? String)?.URLSafeBase64ToBase64(),
            let salt = Data(base64Encoded: saltBase64),
            let serverPublicKeyBase64 = (userInfo[serverPublicKeyUserInfoKey] as? String)?.URLSafeBase64ToBase64(),
            let serverPublicKeyData = Data(base64Encoded: serverPublicKeyBase64)
        else { throw NotificationServiceError.userInfoDataAbsent }

        let secretsService = Secrets(identityId: identityId, keychain: LiveKeychain.self)

        guard
            let auth = try secretsService.getPushAuth(),
            let pushKey = try secretsService.getPushKey()
        else { throw NotificationServiceError.keychainDataAbsent }

        return (try decrypt(encryptedMessage: encryptedMessage,
                           privateKeyData: pushKey,
                           serverPublicKeyData: serverPublicKeyData,
                           auth: auth,
                           salt: salt),
                identityId)
    }

    static func decrypt(encryptedMessage: Data,
                        privateKeyData: Data,
                        serverPublicKeyData: Data,
                        auth: Data,
                        salt: Data) throws -> Data {
        let privateKey = try P256.KeyAgreement.PrivateKey(x963Representation: privateKeyData)
        let serverPublicKey = try P256.KeyAgreement.PublicKey(x963Representation: serverPublicKeyData)
        let sharedSecret = try privateKey.sharedSecretFromKeyAgreement(with: serverPublicKey)

        var keyInfo = HKDFInfo.aesgcm.bytes
        var nonceInfo = HKDFInfo.nonce.bytes
        var context = Array(curve.utf8)
        let publicKeyData = privateKey.publicKey.x963Representation

        context.append(0)
        context.append(0)
        context.append(UInt8(publicKeyData.count))
        context += Array(publicKeyData)
        context.append(0)
        context.append(UInt8(serverPublicKeyData.count))
        context += Array(serverPublicKeyData)

        keyInfo += context
        nonceInfo += context

        let pseudoRandomKey = sharedSecret.hkdfDerivedSymmetricKey(
            using: SHA256.self,
            salt: auth,
            sharedInfo: HKDFInfo.auth.bytes,
            outputByteCount: pseudoRandomKeyLength)
        let key = HKDF<SHA256>.deriveKey(
            inputKeyMaterial: pseudoRandomKey,
            salt: salt,
            info: keyInfo,
            outputByteCount: keyLength)
        let nonce = HKDF<SHA256>.deriveKey(
            inputKeyMaterial: pseudoRandomKey,
            salt: salt,
            info: nonceInfo,
            outputByteCount: nonceLength)

        let sealedBox = try AES.GCM.SealedBox(combined: nonce.withUnsafeBytes(Array.init) + encryptedMessage)
        let decrypted = try AES.GCM.open(sealedBox, using: key)
        let unpadded = decrypted.suffix(from: paddedByteCount)

        return Data(unpadded)
    }

    static func addImage(pushNotification: PushNotification,
                         bestAttemptContent: UNMutableNotificationContent,
                         contentHandler: @escaping (UNNotificationContent) -> Void) {
        let fileName = pushNotification.icon.lastPathComponent
        let fileURL = URL(fileURLWithPath: NSTemporaryDirectory(), isDirectory: true)
            .appendingPathComponent(fileName)

        KingfisherManager.shared.retrieveImage(with: pushNotification.icon) {
            switch $0 {
            case let .success(result):
                let format: ImageFormat

                switch fileURL.pathExtension.lowercased() {
                case "jpg", "jpeg":
                    format = .JPEG
                case "gif":
                    format = .GIF
                case "png":
                    format = .PNG
                default:
                    format = .unknown
                }

                do {
                    try result.image.kf.data(format: format)?.write(to: fileURL)
                    bestAttemptContent.attachments =
                        [try UNNotificationAttachment(identifier: fileName, url: fileURL)]
                    contentHandler(bestAttemptContent)
                } catch {
                    contentHandler(bestAttemptContent)
                }
            case .failure:
                contentHandler(bestAttemptContent)
            }
        }
    }
}

extension String {
    func URLSafeBase64ToBase64() -> String {
        var base64 = replacingOccurrences(of: "-", with: "+").replacingOccurrences(of: "_", with: "/")
        let countMod4 = count % 4

        if countMod4 != 0 {
            base64.append(String(repeating: "=", count: 4 - countMod4))
        }

        return base64
    }
}