Tusker/Tusker/Caching/ImageCache.swift

//
//  ImageCache.swift
//  Tusker
//
//  Created by Shadowfacts on 8/21/18.
//  Copyright © 2018 Shadowfacts. All rights reserved.
//

import UIKit
import Cache

class ImageCache {
    
    static let avatars = ImageCache(name: "Avatars", memoryExpiry: .seconds(60 * 60), diskExpiry: .seconds(60 * 60 * 24))
    static let headers = ImageCache(name: "Headers", memoryExpiry: .seconds(60 * 5), diskExpiry: .seconds(60 * 60))
    static let attachments = ImageCache(name: "Attachments", memoryExpiry: .seconds(60 * 2))
    static let emojis = ImageCache(name: "Emojis", memoryExpiry: .seconds(60 * 5), diskExpiry: .seconds(60 * 60))
    
    #if DEBUG
    private static let disableCaching = ProcessInfo.processInfo.environment.keys.contains("DISABLE_IMAGE_CACHE")
    #else
    private static let disableCaching = false
    #endif

    private let cache: ImageDataCache
    
    private var groups = MultiThreadDictionary<URL, RequestGroup>(name: "ImageCache request groups")
    
    private var backgroundQueue = DispatchQueue(label: "ImageCache completion queue", qos: .default)
    
    init(name: String, memoryExpiry: Expiry, diskExpiry: Expiry? = nil) {
        self.cache = ImageDataCache(name: name, memoryExpiry: memoryExpiry, diskExpiry: diskExpiry)
    }
    
    func get(_ url: URL, completion: ((Data?, UIImage?) -> Void)?) -> Request? {
        let key = url.absoluteString
        if !ImageCache.disableCaching,
           // todo: calling object(forKey: key) does disk I/O and this method is often called from the main thread
           // in performance sensitive paths. a nice optimization to DiskStorage would be adding an internal cache
           // of the state (unknown/exists/does not exist) of whether or not objects exist on disk so that the slow, disk I/O
           // path can be avoided most of the time
           let (data, image) = try? cache.get(key) {
            backgroundQueue.async {
                completion?(data, image)
            }
            return nil
        } else {
            if let completion = completion, let group = groups[url] {
                return group.addCallback(completion)
            } else {
                let group = RequestGroup(url: url) { (data, image) in
                    if let data = data,
                       let image = UIImage(data: data) {
                        try? self.cache.set(key, data: data, image: image)
                    }
                    self.groups.removeValueWithoutReturning(forKey: url)
                }
                groups[url] = group
                let request = group.addCallback(completion)
                group.run()
                return request
            }
        }
    }
    
    func getData(_ url: URL) -> Data? {
        return try? cache.getData(url.absoluteString)
    }
    
    func get(_ url: URL) -> (Data, UIImage)? {
        return try? cache.get(url.absoluteString)
    }

    func cancelWithoutCallback(_ url: URL) {
        groups[url]?.cancelWithoutCallback()
    }
    
    func reset() throws {
        try cache.removeAll()
    }
    
    private class RequestGroup {
        let url: URL
        private let onFinished: (Data?, UIImage?) -> Void
        private var task: URLSessionDataTask?
        private var requests = [Request]()

        init(url: URL, onFinished: @escaping (Data?, UIImage?) -> Void) {
            self.url = url
            self.onFinished = onFinished
        }

        deinit {
            task?.cancel()
        }
        
        func run() {
            task = URLSession.shared.dataTask(with: url, completionHandler: { (data, response, error) in
                guard error == nil, let data = data else {
                    self.complete(with: nil)
                    return
                }
                self.complete(with: data)
            })
            task!.resume()
        }
        
        private func updatePriority() {
            task?.priority = max(1.0, URLSessionTask.defaultPriority + 0.1 * Float(requests.filter { !$0.cancelled }.count))
        }
        
        func addCallback(_ completion: ((Data?, UIImage?) -> Void)?) -> Request {
            let request = Request(callback: completion)
            requests.append(request)
            updatePriority()
            return request
        }
        
        func cancelWithoutCallback() {
            if let request = requests.first(where: { $0.callback == nil && !$0.cancelled }) {
                request.cancel()
                updatePriority()
            }
        }
        
        fileprivate func requestCancelled() {
            let remaining = requests.filter { !$0.cancelled }.count
            if remaining <= 0 {
                task?.cancel()
                complete(with: nil)
            } else {
                updatePriority()
            }
        }

        func complete(with data: Data?) {
            let image = data != nil ? UIImage(data: data!) : nil
            
            requests.filter { !$0.cancelled }.forEach {
                if let callback = $0.callback {
                    callback(data, image)
                }
            }
            
            self.onFinished(data, image)
        }
    }
    
    class Request {
        private weak var group: RequestGroup?
        private(set) var callback: ((Data?, UIImage?) -> Void)?
        private(set) var cancelled: Bool = false
        
        init(callback: ((Data?, UIImage?) -> Void)?) {
            self.callback = callback
        }
        
        func cancel() {
            cancelled = true
            callback = nil
            group?.requestCancelled()
        }
    }
    
}