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?
    
    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 let image {
                    if !loadOriginal,
                       let size = self.desiredPixelSize {
                        image.prepareThumbnail(of: size) {
                            completion(data, $0)
                        }
                    } else {
                        image.prepareForDisplay {
                            completion(data, $0)
                        }
                    }
                } else {
                    completion(data, image)
                }
            }
        } else {
            wrappedCompletion = nil
        }
        
        if !ImageCache.disableCaching,
           let entry = try? cache.get(key, loadOriginal: loadOriginal) {
            wrappedCompletion?(entry.data, entry.image)
            return nil
        } else {
            let task = dataTask(url: url, completion: wrappedCompletion)
            task.resume()
            return task
        }
    }
    
    func get(_ url: URL, loadOriginal: Bool = false) async -> (Data?, UIImage?) {
        // todo: this should integrate with the task cancellation mechanism somehow
        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) {
            let task = dataTask(url: url, completion: nil)
            task.resume()
        }
    }
    
    private func dataTask(url: URL, completion: ((Data?, UIImage?) -> Void)?) -> URLSessionDataTask {
        return URLSession.shared.dataTask(with: url) { data, response, error in
            guard error == nil,
                  let data else {
                return
            }
            let image = UIImage(data: data)
            try? self.cache.set(url.absoluteString, data: data, image: image)
            completion?(data, image)
        }
    }
    
    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 reset() throws {
        try cache.removeAll()
    }
    
    func getDiskSizeInBytes() -> Int64? {
        return cache.disk?.getSizeInBytes()
    }
    
    typealias Request = URLSessionDataTask
    
}