iOS: Rich Notifications

#import "NotificationService.h"

@interface NotificationService ()

@property (nonatomic, strong) void (^contentHandler)(UNNotificationContent *contentToDeliver);
@property (nonatomic, strong) UNMutableNotificationContent *bestAttemptContent;
@property (nonatomic, strong) NSURLSessionDownloadTask *downloadTask;

@end

@implementation NotificationService

- (void)didReceiveNotificationRequest:(UNNotificationRequest *)request withContentHandler:(void (^)(UNNotificationContent * _Nonnull))contentHandler {
    self.contentHandler = contentHandler;
    self.bestAttemptContent = [request.content mutableCopy];
    
    // Modify the notification content here...
    [self marigoldRichNotificationAttachments:self.bestAttemptContent withResponse:^(UNMutableNotificationContent * _Nullable content) {
        self.bestAttemptContent = content;
        self.contentHandler(self.bestAttemptContent);
    }];
}

- (void)serviceExtensionTimeWillExpire {
    // Called just before the extension will be terminated by the system.
    // Use this as an opportunity to deliver your "best attempt" at modified content, otherwise the original push payload will be used.
    [self.downloadTask cancel];
    
    self.contentHandler(self.bestAttemptContent);
}

- (void)marigoldRichNotificationAttachments:(UNMutableNotificationContent *)originalContent withResponse:(nullable void(^)(UNMutableNotificationContent *__nullable modifiedContent))block  {
    // For Image or Video in-app messages, we will send the media URL in the
    // _st payload
    NSString *imageURL = originalContent.userInfo[@"_st"][@"image_url"];
    NSString *videoURL = originalContent.userInfo[@"_st"][@"video_url"];
    
    NSURL *attachmentURL = nil;
    if (videoURL && ![videoURL isKindOfClass:[NSNull class]]) { //Prioritize videos over image
        attachmentURL = [NSURL URLWithString:videoURL];
    }
    else if (imageURL && ![imageURL isKindOfClass:[NSNull class]]) {
        attachmentURL = [NSURL URLWithString:imageURL];
    }
    else {
        block(originalContent); //Nothing to add to the push, return early.
        return;
    }
    
    NSURLSession *session = [NSURLSession sessionWithConfiguration:[NSURLSessionConfiguration defaultSessionConfiguration]];
    self.downloadTask = [session downloadTaskWithURL:attachmentURL completionHandler:^(NSURL *fileLocation, NSURLResponse *response, NSError *error) {
        if (error != nil) {
            block(originalContent); //Nothing to add to the push, return early.
            return;
        }
        else {
            NSFileManager *fileManager = [NSFileManager defaultManager];
            NSString *fileSuffix = attachmentURL.lastPathComponent;
            
            NSURL *typedAttachmentURL = [NSURL fileURLWithPath:[(NSString *_Nonnull)fileLocation.path stringByAppendingString:fileSuffix]];
            [fileManager moveItemAtURL:fileLocation toURL:typedAttachmentURL error:&error];
            
            NSError *attachError = nil;
            UNNotificationAttachment *attachment = [UNNotificationAttachment attachmentWithIdentifier:@"" URL:typedAttachmentURL options:nil error:&attachError];
            
            if (attachment == nil) {
                block(originalContent); //Nothing to add to the push, return early.
                return;
            }
            
            UNMutableNotificationContent *modifiedContent = originalContent.mutableCopy;
            [modifiedContent setAttachments:[NSArray arrayWithObject:attachment]];
            block(modifiedContent);
        }
    }];
    [self.downloadTask resume];
}

@end

import UserNotifications

class NotificationService: UNNotificationServiceExtension {
    
    var contentHandler: ((UNNotificationContent) -> Void)?
    var bestAttemptContent: UNMutableNotificationContent?
    var downloadTask: URLSessionDownloadTask?
    
    override func didReceive(_ request: UNNotificationRequest, withContentHandler contentHandler: @escaping (UNNotificationContent) -> Void) {
        self.contentHandler = contentHandler
        bestAttemptContent = (request.content.mutableCopy() as? UNMutableNotificationContent)
        
        if let bestAttemptContent = bestAttemptContent {
            if let title = bestAttemptContent.userInfo["title"] {
                bestAttemptContent.title = title as! String
            }
            
            var urlString:String?
            
            // Prioritize video over image
            if let stPayload = bestAttemptContent.userInfo["_st"] as? Dictionary<AnyHashable, Any> {
                if let videoURL = stPayload["video_url"] {
                    urlString = videoURL as? String
                } else if let imageURL = stPayload["image_url"] {
                    urlString = imageURL as? String
                } else {
                    // Nothing to add to the push, return early.
                    contentHandler(bestAttemptContent)
                    return
                }
            }
            
            sailthruHandleAttachmentDownload(content: bestAttemptContent.userInfo, urlString: urlString!)
    
        }
    }
    
    func sailthruHandleAttachmentDownload(content: [AnyHashable : Any], urlString: String) {
        
        guard let url = URL(string: urlString) else {
            // Cannot create a valid URL, return early.
            self.contentHandler!(self.bestAttemptContent!)
            return
        }
        
        self.downloadTask = URLSession.shared.downloadTask(with: url) { (location, response, error) in
            if let location = location {
                let tmpDirectory = NSTemporaryDirectory()
                let tmpFile = "file://".appending(tmpDirectory).appending(url.lastPathComponent)
                
                let tmpUrl = URL(string: tmpFile)!
                try! FileManager.default.moveItem(at: location, to: tmpUrl)
                
                if let attachment = try? UNNotificationAttachment(identifier: "", url: tmpUrl) {
                    self.bestAttemptContent?.attachments = [attachment]
                }
            }
            
            self.contentHandler!(self.bestAttemptContent!)
        }
        
        self.downloadTask?.resume()
    }
    
    override func serviceExtensionTimeWillExpire() {
        // Called just before the extension will be terminated by the system.
        // Use this as an opportunity to deliver your "best attempt" at modified content, otherwise the original push payload will be used.
        self.downloadTask?.cancel()
        if let contentHandler = contentHandler, let bestAttemptContent =  bestAttemptContent {
            contentHandler(bestAttemptContent)
        }
    }
}