Caps-iOS/CapCollector/Data/Download.swift

//
//  Download.swift
//  CapCollector
//
//  Created by Christoph on 26.04.20.
//  Copyright © 2020 CH. All rights reserved.
//

import Foundation
import UIKit

final class Download {
    
    let serverUrl: URL
    
    let session: URLSession
    
    let delegate: Delegate
    
    private var downloadingMainImages = Set<Int>()
 
    init(server: URL) {
        let delegate = Delegate()
        
        self.serverUrl = server
        self.session = URLSession(configuration: .ephemeral, delegate: delegate, delegateQueue: nil)
        self.delegate = delegate
    }
    
    // MARK: Paths
    
    var serverNameListUrl: URL {
        Download.serverNameListUrl(server: serverUrl)
    }

    private static func serverNameListUrl(server: URL) -> URL {
        server.appendingPathComponent("names.txt")
    }
    
    private var serverClassifierVersionUrl: URL {
        serverUrl.appendingPathComponent("classifier.version")
    }
    
    private var serverRecognitionModelUrl: URL {
        serverUrl.appendingPathComponent("classifier.mlmodel")
    }
    
    private var serverAllCountsUrl: URL {
        serverUrl.appendingPathComponent("counts")
    }
    
    var serverImageUrl: URL {
        serverUrl.appendingPathComponent("images")
    }
    
    private func serverImageUrl(for cap: Int, version: Int = 0) -> URL {
        serverImageUrl.appendingPathComponent(String(format: "%04d/%04d-%02d.jpg", cap, cap, version))
    }
    
    private func serverNameUrl(for cap: Int) -> URL {
        serverUrl.appendingPathComponent("name/\(cap)")
    }
    
    private func serverImageCountUrl(for cap: Int) -> URL {
        serverUrl.appendingPathComponent("count/\(cap)")
    }
    
    // MARK: Delegate
    
    final class Delegate: NSObject, URLSessionDownloadDelegate {
        
        typealias ProgressHandler = (_ progress: Float, _ bytesWritten: Int64, _ totalBytes: Int64) -> Void
        
        typealias CompletionHandler = (_ url: URL?) -> Void
        
        private var progress = [URLSessionDownloadTask : Float]()
        
        private var callbacks = [URLSessionDownloadTask : ProgressHandler]()
        
        private var completions = [URLSessionDownloadTask : CompletionHandler]()
        
        func registerForProgress(_ downloadTask: URLSessionDownloadTask, callback: ProgressHandler?, completion: @escaping CompletionHandler) {
            progress[downloadTask] = 0
            callbacks[downloadTask] = callback
            completions[downloadTask] = completion
        }
        
        func urlSession(_ session: URLSession, downloadTask: URLSessionDownloadTask, didFinishDownloadingTo location: URL) {
            completions[downloadTask]?(location)
            callbacks[downloadTask] = nil
            progress[downloadTask] = nil
            completions[downloadTask] = nil
        }
        
        func urlSession(_ session: URLSession, downloadTask: URLSessionDownloadTask, didWriteData bytesWritten: Int64, totalBytesWritten: Int64, totalBytesExpectedToWrite: Int64) {
            let ratio = totalBytesExpectedToWrite > 0 ? Float(totalBytesWritten) / Float(totalBytesExpectedToWrite) : 0
            progress[downloadTask] = ratio
            callbacks[downloadTask]?(ratio, totalBytesWritten, totalBytesExpectedToWrite)
        }
        
        func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
            guard let downloadTask = task as? URLSessionDownloadTask else {
                return
            }
            completions[downloadTask]?(nil)
            callbacks[downloadTask] = nil
            progress[downloadTask] = nil
            completions[downloadTask] = nil
        }
    }
    
    // MARK: Downloading data
    
    func image(for cap: Int, to url: URL, timeout: TimeInterval = 30) -> Bool {
        let group = DispatchGroup()
        group.enter()
        var result = true
        let success = image(for: cap, version: 0, to: url) { success in
            result = success
            group.leave()
        }
        guard success else {
            log("Already downloading image for cap \(cap)")
            return false
        }
        guard group.wait(timeout: .now() + timeout) == .success else {
            log("Timed out downloading image for cap \(cap)")
            return false
        }
        return result
    }
    
    /**
     Download an image for a cap.
     - Parameter cap: The id of the cap.
     - Parameter version: The image version to download.
     - Parameter completion: A closure with the resulting image
     - Returns: `true`, of the file download was started, `false`, if the image is already downloading.
     */
    @discardableResult
    func image(for cap: Int, version: Int = 0, to url: URL, queue: DispatchQueue = .main, completion: @escaping (Bool) -> Void) -> Bool {
        // Check if main image, and already being downloaded
        if version == 0 {
            guard !downloadingMainImages.contains(cap) else {
                return false
            }
            downloadingMainImages.insert(cap)
        }
        let serverUrl = serverImageUrl(for: cap, version: version)
        let query = "Image of cap \(cap) version \(version)"
        let task = session.downloadTask(with: serverUrl) { fileUrl, response, error in
            if version == 0 {
                queue.async {
                    self.downloadingMainImages.remove(cap)
                }
            }
            guard let fileUrl = self.convertResponse(to: query, fileUrl, response, error) else {
                completion(false)
                return
            }
            do {
                if FileManager.default.fileExists(atPath: url.path) {
                    try FileManager.default.removeItem(at: url)
                }
                try FileManager.default.moveItem(at: fileUrl, to: url)
            } catch {
                self.log("Failed to move downloaded image for cap \(cap): \(error)")
                completion(false)
            }
            completion(true)
        }
        task.resume()
        return true
    }
    
    func imageCount(for cap: Int, completion: @escaping (_ count: Int?) -> Void) {
        let url = serverImageCountUrl(for: cap)
        let query = "Image count for cap \(cap)"
        session.startTaskExpectingInt(with: url, query: query, completion: completion)
    }
    
    func name(for cap: Int, completion: @escaping (_ name: String?) -> Void) {
        let url = serverNameUrl(for: cap)
        let query = "Name for cap \(cap)"
        session.startTaskExpectingString(with: url, query: query, completion: completion)
    }
    
    func imageCounts(completion: @escaping ([Int]?) -> Void) {
        let query = "Image count of all caps"
        session.startTaskExpectingData(with: serverAllCountsUrl, query: query) { data in
            guard let data = data else {
                completion(nil)
                return
            }
            completion(data.map(Int.init))
        }
    }
    
    func names(completion: @escaping ([String]?) -> Void) {
        let query = "Download of server database"
        session.startTaskExpectingString(with: serverNameListUrl, query: query) { string in
            completion(string?.trimmingCharacters(in: .whitespacesAndNewlines).components(separatedBy: "\n"))
        }
    }
    
    func databaseSize(completion: @escaping (_ size: Int64?) -> Void) {
        size(of: "database size", to: serverNameListUrl, completion: completion)
    }
    
    func classifierVersion(completion: @escaping (Int?) -> Void) {
        let query = "Server classifier version"
        session.startTaskExpectingInt(with: serverClassifierVersionUrl, query: query, completion: completion)
    }
    
    func classifierSize(completion: @escaping (Int64?) -> Void) {
        size(of: "classifier size", to: serverRecognitionModelUrl, completion: completion)
    }
    
    func classifier(progress: Delegate.ProgressHandler? = nil, completion: @escaping (URL?) -> Void) {
        let task = session.downloadTask(with: serverRecognitionModelUrl)
        delegate.registerForProgress(task, callback: progress) { url in
            self.log("Classifier download complete")
            completion(url)
        }
        task.resume()
    }
    
    // MARK: Requests
    
    private func size(of query: String, to url: URL, completion: @escaping (_ size: Int64?) -> Void) {
        var request = URLRequest(url: url)
        request.httpMethod = "HEAD"
        let task = session.dataTask(with: request)  { _, response, _ in
            guard let r = response else {
                self.log("Request '\(query)' received no response")
                completion(nil)
                return
            }
            completion(r.expectedContentLength)
        }
        task.resume()
    }
    
    private func convertIntResponse(to query: String, _ data: Data?, _ response: URLResponse?, _ error: Error?) -> Int? {
        guard let string = self.convertStringResponse(to: query, data, response, error) else {
            return nil
        }
        guard let int = Int(string.trimmingCharacters(in: .whitespacesAndNewlines)) else {
            self.log("Request '\(query)' received an invalid value '\(string)'")
            return nil
        }
        return int
    }
    
    private func convertStringResponse(to query: String, _ data: Data?, _ response: URLResponse?, _ error: Error?) -> String? {
        guard let data = self.convertResponse(to: query, data, response, error) else {
            return nil
        }
        guard let string = String(data: data, encoding: .utf8) else {
            self.log("Request '\(query)' received invalid data (not a string)")
            return nil
        }
        return string
    }
    
    private func convertResponse<T>(to query: String, _ result: T?, _ response: URLResponse?, _ error: Error?) -> T? {
        if let error = error {
            log("Request '\(query)' produced an error: \(error)")
            return nil
        }
        guard let response = response else {
            log("Request '\(query)' received no response")
            return nil
        }
        guard let urlResponse = response as? HTTPURLResponse else {
            log("Request '\(query)' received an invalid response: \(response)")
            return nil
        }
        guard urlResponse.statusCode == 200 else {
            log("Request '\(query)' failed with status code \(urlResponse.statusCode)")
            return nil
        }
        guard let r = result else {
            log("Request '\(query)' received no item")
            return nil
        }
        return r
    }
    
    
}

extension Download: Logger { }

extension URLSession {
    
    func startTaskExpectingData(with url: URL, query: String, completion: @escaping (Data?) -> Void) {
        let task = dataTask(with: url) { data, response, error in
            if let error = error {
                log("Request '\(query)' produced an error: \(error)")
                completion(nil)
                return
            }
            guard let response = response else {
                log("Request '\(query)' received no response")
                completion(nil)
                return
            }
            guard let urlResponse = response as? HTTPURLResponse else {
                log("Request '\(query)' received an invalid response: \(response)")
                completion(nil)
                return
            }
            guard urlResponse.statusCode == 200 else {
                log("Request '\(query)' failed with status code \(urlResponse.statusCode)")
                completion(nil)
                return
            }
            guard let d = data else {
                log("Request '\(query)' received no data")
                completion(nil)
                return
            }
            completion(d)
        }
        task.resume()
    }
    
    func startTaskExpectingString(with url: URL, query: String, completion: @escaping (String?) -> Void) {
        startTaskExpectingData(with: url, query: query) { data in
            guard let data = data else {
                completion(nil)
                return
            }
            guard let string = String(data: data, encoding: .utf8) else {
                log("Request '\(query)' received invalid data (not a string)")
                completion(nil)
                return
            }
            completion(string)
        }
    }
    
    func startTaskExpectingInt(with url: URL, query: String, completion: @escaping (Int?) -> Void) {
        startTaskExpectingString(with: url, query: query) { string in
            guard let string = string else {
                completion(nil)
                return
            }
            guard let int = Int(string.trimmingCharacters(in: .whitespacesAndNewlines)) else {
                log("Request '\(query)' received an invalid value '\(string)'")
                completion(nil)
                return
            }
            completion(int)
        }
    }
}