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

    private let cache: Cache<Data>
    
    private var groups = [URL: RequestGroup]()
    
    init(name: String, memoryExpiry expiry: Expiry) {
        let storage = MemoryStorage<Data>(config: MemoryConfig(expiry: expiry))
        self.cache = .memory(storage)
    }
    
    init(name: String, diskExpiry expiry: Expiry) {
        let storage = try! DiskStorage<Data>(config: DiskConfig(name: name, expiry: expiry), transformer: TransformerFactory.forData())
        self.cache = .disk(storage)
    }
    
    init(name: String, memoryExpiry: Expiry, diskExpiry: Expiry) {
        let memory = MemoryStorage<Data>(config: MemoryConfig(expiry: memoryExpiry))
        let disk = try! DiskStorage<Data>(config: DiskConfig(name: name, expiry: diskExpiry), transformer: TransformerFactory.forData())
        self.cache = .hybrid(HybridStorage(memoryStorage: memory, diskStorage: disk))
    }
    
    func get(_ url: URL, completion: ((Data?) -> Void)?) -> Request? {
        let key = url.absoluteString
        if let data = try? cache.object(forKey: key) {
            completion?(data)
            return nil
        } else {
            if let completion = completion, let group = groups[url] {
                return group.addCallback(completion)
            } else {
                let group = RequestGroup(url: url) { (data) in
                    if let data = data {
                        try? self.cache.setObject(data, forKey: key)
                    }
                    self.groups.removeValue(forKey: url)
                }
                groups[url] = group
                let request = group.addCallback(completion)
                group.run()
                return request
            }
        }
    }
    
    func get(_ url: URL) -> Data? {
        return try? cache.object(forKey: 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?) -> Void
        private var task: URLSessionDataTask?
        private var requests = [Request]()

        init(url: URL, onFinished: @escaping (Data?) -> 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?) -> 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?) {
            requests.filter { !$0.cancelled }.forEach {
                if let callback = $0.callback {
                    callback(data)
                }
            }
            self.onFinished(data)
        }
    }
    
    class Request {
        private weak var group: RequestGroup?
        private(set) var callback: ((Data?) -> Void)?
        private(set) var cancelled: Bool = false
        
        init(callback: ((Data?) -> Void)?) {
            self.callback = callback
        }
        
        func cancel() {
            cancelled = true
            callback = nil
            group?.requestCancelled()
        }
    }
    
}