Tusker/Tusker/Caching/ImageCache.swift

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

import UIKit

class ImageCache {
    
    static let avatars = ImageCache(name: "Avatars", memoryExpiry: .seconds(60 * 60), diskExpiry: .seconds(60 * 60 * 24 * 7), desiredSize: CGSize(width: 50, height: 50))
    static let headers = ImageCache(name: "Headers", memoryExpiry: .seconds(60 * 5), diskExpiry: .seconds(60 * 60 * 24 * 7))
    static let attachments = ImageCache(name: "Attachments", memoryExpiry: .seconds(60 * 2))
    static let emojis = ImageCache(name: "Emojis", memoryExpiry: .seconds(60 * 5), diskExpiry: .seconds(60 * 60 * 24 * 7))
    
    #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 let desiredPixelSize: CGSize?
    
    private var groups = MultiThreadDictionary<URL, RequestGroup>()
    
    private var backgroundQueue = DispatchQueue(label: "ImageCache completion queue", qos: .default)
    
    init(name: String, memoryExpiry: CacheExpiry, diskExpiry: CacheExpiry? = nil, desiredSize: CGSize? = nil) {
        // todo: might not always want to use UIScreen.main for this, e.g. Catalyst?
        let pixelSize = desiredSize?.applying(.init(scaleX: UIScreen.main.scale, y: UIScreen.main.scale))
        self.desiredPixelSize = pixelSize
        self.cache = ImageDataCache(name: name, memoryExpiry: memoryExpiry, diskExpiry: diskExpiry, storeOriginalDataInMemory: diskExpiry == nil, desiredPixelSize: pixelSize)
    }
    
    func get(_ url: URL, loadOriginal: Bool = false, completion: ((Data?, UIImage?) -> Void)?) -> Request? {
        let key = url.absoluteString
        
        let wrappedCompletion: ((Data?, UIImage?) -> Void)?
        if let completion = completion {
            wrappedCompletion = { (data, image) in
                if !loadOriginal,
                   let size = self.desiredPixelSize {
                    image?.prepareThumbnail(of: size, completionHandler: {
                        completion(data, $0)
                    })
                } else {
                    image?.prepareForDisplay {
                        completion(data, $0)
                    }
                }
            }
        } else {
            wrappedCompletion = nil
        }
        
        if !ImageCache.disableCaching,
           let entry = try? cache.get(key, loadOriginal: loadOriginal) {
            wrappedCompletion?(entry.data, entry.image)
            return nil
        } else {
            if let group = groups[url] {
                return group.addCallback(wrappedCompletion)
            } else {
                let group = createGroup(url: url)
                let request = group.addCallback(wrappedCompletion)
                group.run()
                return request
            }
        }
    }
    
    func get(_ url: URL, loadOriginal: Bool = false) async -> (Data?, UIImage?) {
        return await withCheckedContinuation { continuation in
            _ = get(url, loadOriginal: loadOriginal) { data, image in
                continuation.resume(returning: (data, image))
            }
        }
    }
    
    func fetchIfNotCached(_ url: URL) {
        // if caching is disabled, don't bother fetching since nothing will be done with the result
        guard !ImageCache.disableCaching else { return }
        
        if !((try? cache.has(url.absoluteString)) ?? false),
           !groups.contains(key: url) {
            let group = createGroup(url: url)
            group.run()
        }
    }
    
    private func createGroup(url: URL) -> RequestGroup {
        let group = RequestGroup(url: url) { (data, image) in
            if let data = data {
                try? self.cache.set(url.absoluteString, data: data, image: image)
            }
            _ = self.groups.removeValue(forKey: url)
        }
        groups[url] = group
        return group
    }
    
    func getData(_ url: URL) -> Data? {
        return try? cache.getData(url.absoluteString)
    }
    
    func get(_ url: URL, loadOriginal: Bool = false) -> ImageDataCache.Entry? {
        return try? cache.get(url.absoluteString, loadOriginal: loadOriginal)
    }

    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()
        }
        
        func addCallback(_ completion: ((Data?, UIImage?) -> Void)?) -> Request {
            let request = Request(callback: completion)
            requests.append(request)
            return request
        }
        
        func cancelWithoutCallback() {
            if let request = requests.first(where: { $0.callback == nil && !$0.cancelled }) {
                request.cancel()
            }
        }
        
        fileprivate func requestCancelled() {
            let remaining = requests.filter { !$0.cancelled }.count
            if remaining <= 0 {
                task?.cancel()
                complete(with: nil)
            }
        }

        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()
        }
    }
    
}