Caps-iOS/CapCollector/Data/Storage.swift

//
//  DiskManager.swift
//  CapFinder
//
//  Created by User on 23.04.18.
//  Copyright © 2018 User. All rights reserved.
//

import Foundation
import UIKit
import CoreML
import Vision

final class Storage {
    
    // MARK: Paths
    
    let fm = FileManager.default
    
    let baseUrl: URL
    
    // MARK: INIT
    
    init(in folder: URL) {
        self.baseUrl = folder
    }
    
    // MARK: File/folder urls
    
    var dbUrl: URL {
        baseUrl.appendingPathComponent("db.sqlite3")
    }
    
    var modelUrl: URL {
        baseUrl.appendingPathComponent("model.mlmodel")
    }
    
    private func localImageUrl(for cap: Int, version: Int) -> URL {
        baseUrl.appendingPathComponent("\(cap)-\(version).jpg")
    }
    
    private func thumbnailUrl(for cap: Int) -> URL {
        baseUrl.appendingPathComponent("\(cap)-thumb.jpg")
    }
    
    // MARK: Storage
    
    /**
     Save an image to disk
     - parameter url: The url where the downloaded image is stored
     - parameter cap: The cap id
     - parameter version: The version of the image to get
     - returns: True, if the image was saved
    */
    func saveImage(at url: URL, for cap: Int, version: Int = 0) -> Bool {
        let targetUrl = localImageUrl(for: cap, version: version)
        do {
            if fm.fileExists(atPath: targetUrl.path) {
                try fm.removeItem(at: targetUrl)
            }
            try fm.moveItem(at: url, to: targetUrl)
            return true
        } catch {
            log("Failed to delete or move image \(version) for cap \(cap)")
            return false
        }
    }
    
    /**
     Save an image to disk
     - parameter image: The image
     - parameter cap: The cap id
     - parameter version: The version of the image
     - returns: True, if the image was saved
    */
    func save(image: UIImage, for cap: Int, version: Int = 0) -> Bool {
        guard let data = image.jpegData(compressionQuality: Cap.jpgQuality) else {
            return false
        }
        return save(imageData: data, for: cap, version: version)
    }
    
    /**
     Save image data to disk
     - parameter image: The data of the image
     - parameter cap: The cap id
     - parameter version: The version of the image
     - returns: True, if the image was saved
    */
    func save(imageData: Data, for cap: Int, version: Int = 0) -> Bool {
        write(imageData, to: localImageUrl(for: cap, version: version))
    }
    
    /**
     Save a thumbnail.
     - parameter thumbnail: The image
     - parameter cap: The cap id
     - returns: True, if the image was saved
     */
    func save(thumbnail: UIImage, for cap: Int) -> Bool {
        guard let data = thumbnail.jpegData(compressionQuality: Cap.jpgQuality) else {
            return false
        }
        return save(thumbnailData: data, for: cap)
    }
    
    /**
     Save a thumbnail to the download folder
     - parameter thumbnailData: The data of the image
     - parameter cap: The cap id
     - returns: True, if the image was saved
     */
    func save(thumbnailData: Data, for cap: Int) -> Bool {
        write(thumbnailData, to: thumbnailUrl(for: cap))
    }
    
    /**
     Save the downloaded and compiled recognition model.
     - Parameter url: The temporary location to which the model was compiled.
     - Returns: `true`, if the model was moved.
     */
    func save(recognitionModelAt url: URL) -> Bool {
        move(url, to: modelUrl)
    }
    
    /**
    Save the downloaded and database.
    - Parameter url: The temporary location to which the database was downloaded.
    - Returns: `true`, if the database was moved.
    */
    func save(databaseAt url: URL) -> Bool {
        move(url, to: dbUrl)
    }
    
    private func move(_ url: URL, to destination: URL) -> Bool {
        if fm.fileExists(atPath: destination.path) {
            do {
                try fm.removeItem(at: destination)
            } catch {
                log("Failed to remove file \(destination.lastPathComponent) before writing new version: \(error)")
                return false
            }
        }
        do {
            try fm.moveItem(at: url, to: destination)
            return true
        } catch {
            self.error("Failed to move file \(destination.lastPathComponent) to permanent location: \(error)")
            return false
        }
    }

    private func write(_ data: Data, to url: URL) -> Bool {
        do {
            try data.write(to: url)
        } catch {
            self.error("Could not write data to \(url): \(error)")
            return false
        }
        return true
    }
    
    // MARK: Status
    
    /**
     Check if an image exists for a cap
     - parameter cap: The id of the cap
     - returns: True, if an image exists
     */
    func hasImage(for cap: Int) -> Bool {
        fm.fileExists(atPath: localImageUrl(for: cap, version: 0).path)
    }
    
    func existingImageUrl(for cap: Int, version: Int = 0) -> URL? {
        let url = localImageUrl(for: cap, version: version)
        return fm.fileExists(atPath: url.path) ? url : nil
    }
    
    // MARK: Retrieval
    
    /**
     Get the image data for a cap.
     If the image exists on disk, it is returned.
     If no image exists locally, then this function returns nil.
     - parameter cap: The id of the cap
     - parameter version: The image version
     - returns: The image data, or `nil`
     */
    func imageData(for cap: Int, version: Int = 0) -> Data? {
        readData(from: localImageUrl(for: cap, version: version))
    }
    
    /**
     Get the image for a cap.
     If the image exists on disk, it is returned.
     If no image exists locally, then this function returns nil.
     - parameter cap: The id of the cap
     - parameter version: The image version
     - returns: The image, or `nil`
     - note: Removes invalid image data on disk, if the data is not a valid image
     - note: Must be called on the main thread
     */
    func image(for cap: Int, version: Int = 0) -> UIImage? {
        guard let data = imageData(for: cap, version: version) else {
            return nil
        }
        guard let image = UIImage(data: data) else {
            log("Removing invalid image \(version) of cap \(cap) from disk")
            deleteImage(for: cap, version: version)
            return nil
        }
        return image
    }
    
    /**
     Get the image for a cap.
     If the image exists on disk, it is returned.
     If no image exists locally, then this function returns nil.
     - parameter cap: The id of the cap
     - parameter version: The image version
     - returns: The image, or `nil`
     */
    func ciImage(for cap: Int, version: Int = 0) -> CIImage? {
        guard let url = existingImageUrl(for: cap, version: version) else {
            return nil
        }
        return CIImage(contentsOf: url)
    }
    
    /**
     Get the thumbnail data for a cap.
     If the image exists on disk, it is returned.
     If no image exists locally, then this function returns nil.
     - parameter cap: The id of the cap
     - returns: The image data, or `nil`
     */
    func thumbnailData(for cap: Int) -> Data? {
        readData(from: thumbnailUrl(for: cap))
    }
    
    /**
     Get the thumbnail for a cap.
     If the image exists on disk, it is returned.
     If no image exists locally, then this function returns nil.
     - parameter cap: The id of the cap
     - returns: The image, or `nil`
     */
    func thumbnail(for cap: Int) -> UIImage? {
        guard let data = thumbnailData(for: cap) else {
            return nil
        }
        return UIImage(data: data)
    }
    
    /// The compiled recognition model on disk
    var recognitionModel: VNCoreMLModel? {
        guard fm.fileExists(atPath: modelUrl.path) else {
            log("No recognition model to load")
            return nil
        }
        do {
            let model = try MLModel(contentsOf: modelUrl)
            
            return try VNCoreMLModel(for: model)
        } catch {
            self.error("Failed to load recognition model: \(error)")
            return nil
        }
    }
    
    func averageColor(for cap: Int, version: Int = 0) -> UIColor? {
        guard let inputImage = ciImage(for: cap, version: version) else {
            error("Failed to read CIImage for main image of cap \(cap)")
            return nil
        }
        return inputImage.averageColor
    }
    
    
    private func readData(from url: URL) -> Data? {
        guard fm.fileExists(atPath: url.path) else {
            return nil
        }
        do {
            return try Data(contentsOf: url)
        } catch {
            self.error("Could not read data from \(url): \(error)")
            return nil
        }
    }
    
    // MARK: Deleting data
    
    @discardableResult
    func deleteDatabase() -> Bool {
        do {
            try fm.removeItem(at: dbUrl)
            return true
        } catch {
            log("Failed to delete database: \(error)")
            return false
        }
    }
    
    @discardableResult
    func deleteImage(for cap: Int, version: Int) -> Bool {
        guard let url = existingImageUrl(for: cap, version: version) else {
            return true
        }
        do {
            try fm.removeItem(at: url)
            return true
        } catch {
            log("Failed to delete image \(version) for cap \(cap): \(error)")
            return false
        }
    }
}

extension Storage: Logger { }