// // Database.swift // CapCollector // // Created by Christoph on 14.04.20. // Copyright © 2020 CH. All rights reserved. // import Foundation import UIKit import CoreML import SQLite protocol DatabaseDelegate: AnyObject { func database(didAddCap cap: Cap) func database(didChangeCap cap: Int) func database(didLoadImageForCap cap: Int) func database(completedBackgroundWorkItem title: String, subtitle: String) func database(needsUserConfirmation title: String, body: String, shouldProceed: @escaping (Bool) -> Void) func database(didFailBackgroundWork title: String, subtitle: String) func databaseHasNewClassifier() func databaseDidFinishBackgroundWork() func databaseNeedsFullRefresh() } private enum BackgroundWorkTaskType: Int, CustomStringConvertible, Comparable { case downloadCapNames = 9 case downloadCounts = 8 case downloadClassifier = 7 case uploadingCaps = 6 case uploadingImages = 5 case downloadMainImages = 4 case creatingThumbnails = 3 case creatingColors = 2 var description: String { switch self { case .downloadCapNames: return "Downloading names" case .downloadCounts: return "Downloading counts" case .downloadClassifier: return "Downloading classifier" case .uploadingCaps: return "Uploading caps" case .uploadingImages: return "Uploading images" case .downloadMainImages: return "Downloading images" case .creatingThumbnails: return "Creating thumbnails" case .creatingColors: return "Creating colors" } } var maximumNumberOfSimultaneousItems: Int { switch self { case .downloadMainImages: return 50 case .creatingThumbnails: return 10 case .creatingColors: return 10 default: return 1 } } var nextType: BackgroundWorkTaskType? { BackgroundWorkTaskType(rawValue: rawValue - 1) } static func < (lhs: BackgroundWorkTaskType, rhs: BackgroundWorkTaskType) -> Bool { lhs.rawValue < rhs.rawValue } } final class Database { // MARK: Variables let db: Connection private let upload: Upload private let download: Download let storage: Storage weak var delegate: DatabaseDelegate? init?(url: URL, server: URL, storageFolder: URL) { guard let db = try? Connection(url.path) else { return nil } let upload = Upload(server: server) let download = Download(server: server) do { try db.run(Cap.createQuery) try db.run(upload.createQuery) try db.run(Database.Colors.createQuery) try db.run(Database.TileImage.createQuery) } catch { return nil } self.db = db self.upload = upload self.download = download self.storage = Storage(in: storageFolder) log("Database loaded with \(capCount) caps") } // MARK: Computed properties /// All caps currently in the database var caps: [Cap] { (try? db.prepare(Cap.table))?.map(Cap.init) ?? [] } /// The ids of all caps var capIds: Set { Set(caps.map { $0.id }) } /// A dictionary of all caps, indexed by their ids var capDict: [Int : Cap] { caps.reduce(into: [:]) { $0[$1.id] = $1 } } /// The ids of the caps which weren't included in the last classification var unmatchedCaps: [Int] { let query = Cap.table.select(Cap.columnId).filter(Cap.columnMatched == false) return (try? db.prepare(query).map { $0[Cap.columnId] }) ?? [] } /// The number of caps which could be recognized during the last classification var recognizedCapCount: Int { (try? db.scalar(Cap.table.filter(Cap.columnMatched == true).count)) ?? 0 } /// The number of caps currently in the database var capCount: Int { (try? db.scalar(Cap.table.count)) ?? 0 } /// The total number of images for all caps var imageCount: Int { (try? db.prepare(Cap.table).reduce(0) { $0 + $1[Cap.columnCount] }) ?? 0 } var nextPendingCapUpload: Cap? { do { guard let row = try db.pluck(Cap.table.filter(Cap.columnUploaded == false).order(Cap.columnId.asc)) else { return nil } return Cap(row: row) } catch { log("Failed to get next pending cap upload") return nil } } var pendingCapUploadCount: Int { do { let query = Cap.table.filter(Cap.columnUploaded == false).count return try db.scalar(query) } catch { log("Failed to get pending cap upload count") return 0 } } var nextPendingImageUpload: (id: Int, version: Int)? { do { guard let row = try db.pluck(upload.table) else { return nil } return (id: row[upload.rowCapId], version: row[upload.rowCapVersion]) } catch { log("Failed to get pending image uploads") return nil } } var capsWithImages: Set { capIds.filter { storage.hasImage(for: $0) } } var capsWithThumbnails: Set { capIds.filter { storage.hasThumbnail(for: $0) } } var pendingImageUploadCount: Int { ((try? db.scalar(upload.table.count)) ?? 0) } /// The number of caps without a thumbnail on disk var pendingCapForThumbnailCreation: Int { caps.reduce(0) { $0 + (storage.hasThumbnail(for: $1.id) ? 0 : 1) } } var pendingCapsForColorCreation: Int { do { return try capCount - db.scalar(Colors.table.count) } catch { log("Failed to get count of caps without color: \(error)") return 0 } } var classifierVersion: Int { set { UserDefaults.standard.set(newValue, forKey: Classifier.userDefaultsKey) log("Classifier version set to \(newValue)") } get { UserDefaults.standard.integer(forKey: Classifier.userDefaultsKey) } } var isInOfflineMode: Bool { set { UserDefaults.standard.set(newValue, forKey: Upload.offlineKey) log("Offline mode set to \(newValue)") } get { UserDefaults.standard.bool(forKey: Upload.offlineKey) } } // MARK: Data updates /** Create a new cap with an image. The cap is inserted into the database, and the name and image will be uploaded to the server. - parameter image: The main image of the cap - parameter name: The name of the cap - note: Must be called on the main queue. - note: The registered delegates will be informed about the added cap through `database(didAddCap:)` - returns: `true`, if the cap was created. */ func createCap(image: UIImage, name: String) -> Bool { let cap = Cap(name: name, id: capCount + 1) guard insert(cap: cap) else { log("Cap not inserted") return false } guard storage.save(image: image, for: cap.id) else { log("Cap image not saved") return false } addPendingUpload(for: cap.id, version: 0) startBackgroundWork() return true } /** Insert a new cap. Only inserts the cap into the database, and optionally notifies the delegates. - note: When a new cap is created, use `createCap(image:name:)` instead */ @discardableResult private func insert(cap: Cap, notify: Bool = true) -> Bool { do { try db.run(cap.insertQuery) if notify { DispatchQueue.main.async { self.delegate?.database(didAddCap: cap) } } return true } catch { log("Failed to insert cap \(cap.id): \(error)") return false } } func add(image: UIImage, for cap: Int) -> Bool { guard let version = count(for: cap) else { log("Failed to get count for cap \(cap)") return false } guard storage.save(image: image, for: cap, version: version) else { log("Failed to save image \(version) for cap \(cap) to disk") return false } guard update(count: version + 1, for: cap) else { log("Failed update count \(version) for cap \(cap)") return false } guard addPendingUpload(for: cap, version: version) else { log("Failed to add cap \(cap) version \(version) to upload queue") return false } startBackgroundWork() return true } // MARK: Updating cap properties private func update(_ property: String, for cap: Int, notify: Bool = true, setter: Setter...) -> Bool { do { let query = updateQuery(for: cap).update(setter) try db.run(query) if notify { DispatchQueue.main.async { self.delegate?.database(didChangeCap: cap) } } return true } catch { log("Failed to update \(property) for cap \(cap): \(error)") return false } } @discardableResult private func update(uploaded: Bool, for cap: Int) -> Bool { update("uploaded", for: cap, setter: Cap.columnUploaded <- uploaded) } @discardableResult private func update(count: Int, for cap: Int) -> Bool { update("count", for: cap, setter: Cap.columnCount <- count) } @discardableResult private func update(matched: Bool, for cap: Int) -> Bool { update("matched", for: cap, setter: Cap.columnMatched <- matched) } // MARK: External editing /** Update the `name` of a cap. */ @discardableResult func update(name: String, for cap: Int) -> Bool { guard update("name", for: cap, setter: Cap.columnName <- name, Cap.columnUploaded <- false) else { return false } startBackgroundWork() return true } @discardableResult private func updateWithoutUpload(name: String, for cap: Int) -> Bool { update("name", for: cap, notify: false, setter: Cap.columnName <- name) } func update(recognizedCaps: Set) { let unrecognized = self.unmatchedCaps // Update caps which haven't been recognized before let newlyRecognized = recognizedCaps.intersection(unrecognized) let logIndividualMessages = newlyRecognized.count < 10 if !logIndividualMessages { log("Marking \(newlyRecognized.count) caps as matched") } for cap in newlyRecognized { if logIndividualMessages { log("Marking cap \(cap) as matched") } update(matched: true, for: cap) } // Update caps which are no longer recognized let missing = Set(1...capCount).subtracting(recognizedCaps).subtracting(unrecognized) for cap in missing { log("Marking cap \(cap) as not matched") update(matched: false, for: cap) } } // MARK: Uploads @discardableResult private func addPendingUpload(for cap: Int, version: Int) -> Bool { do { guard try db.scalar(upload.existsQuery(for: cap, version: version)) == 0 else { return true } try db.run(upload.insertQuery(for: cap, version: version)) return true } catch { log("Failed to add pending upload of cap \(cap) version \(version): \(error)") return false } } @discardableResult private func removePendingUpload(for cap: Int, version: Int) -> Bool { do { try db.run(upload.deleteQuery(for: cap, version: version)) return true } catch { log("Failed to remove pending upload of cap \(cap) version \(version): \(error)") return false } } // MARK: Information retrieval func cap(for id: Int) -> Cap? { do { guard let row = try db.pluck(updateQuery(for: id)) else { log("No cap with id \(id) in database") return nil } return Cap(row: row) } catch { log("Failed to get cap \(id): \(error)") return nil } } private func count(for cap: Int) -> Int? { do { let row = try db.pluck(updateQuery(for: cap).select(Cap.columnCount)) return row?[Cap.columnCount] } catch { log("Failed to get count for cap \(cap)") return nil } } func countOfCaps(withImageCountLessThan limit: Int) -> Int { do { return try db.scalar(Cap.table.filter(Cap.columnCount < limit).count) } catch { log("Failed to get caps with less than \(limit) images") return 0 } } func lowestImageCountForCaps(startingAt start: Int) -> (count: Int, numberOfCaps: Int) { do { var currentCount = start - 1 var capsFound = 0 repeat { currentCount += 1 capsFound = try db.scalar(Cap.table.filter(Cap.columnCount == currentCount).count) } while capsFound == 0 return (currentCount, capsFound) } catch { return (0,0) } } func updateQuery(for cap: Int) -> Table { Cap.table.filter(Cap.columnId == cap) } // MARK: Downloads @discardableResult func downloadImage(for cap: Int, version: Int = 0, completion: @escaping (_ image: UIImage?) -> Void) -> Bool { let url = storage.localImageUrl(for: cap, version: version) return download.image(for: cap, version: version, to: url) { success in if version == 0 && success { DispatchQueue.main.async { self.delegate?.database(didLoadImageForCap: cap) } } let image = self.storage.image(for: cap, version: version) completion(image) } } private func update(names: [String]) { let notify = capCount > 0 log("Downloaded cap names (initialDownload: \(!notify))") let caps = self.capDict let changed: [Int] = names.enumerated().compactMap { id, name in let id = id + 1 guard let existingName = caps[id]?.name else { // Insert cap let cap = Cap(id: id, name: name, count: 0) guard insert(cap: cap, notify: notify) else { return nil } return id } guard existingName != name else { // Name unchanged return nil } guard updateWithoutUpload(name: name, for: id) else { return nil } return id } if !notify { log("Added \(changed.count) new caps after initial download") delegate?.databaseNeedsFullRefresh() } } var isDoingWorkInBackgound: Bool { backgroundTaskStatus != nil } private var didUpdateBackgroundItems = false private var backgroundTaskStatus: BackgroundWorkTaskType? = nil private var expectedBackgroundWorkStatus: BackgroundWorkTaskType? = nil private var nextBackgroundWorkStatus: BackgroundWorkTaskType? { guard let oldType = backgroundTaskStatus else { return expectedBackgroundWorkStatus } guard let type = expectedBackgroundWorkStatus else { return backgroundTaskStatus?.nextType } guard oldType > type else { return type } return oldType.nextType } private func setNextBackgroundWorkStatus() -> BackgroundWorkTaskType? { backgroundTaskStatus = nextBackgroundWorkStatus expectedBackgroundWorkStatus = nil return backgroundTaskStatus } private let context = CIContext(options: [.workingColorSpace: kCFNull!]) func startInitialDownload() { startBackgroundWork(startingWith: .downloadCapNames) } func scheduleClassifierDownload() { startBackgroundWork(startingWith: .downloadClassifier) } func startBackgroundWork() { startBackgroundWork(startingWith: .uploadingCaps) } private func startBackgroundWork(startingWith type: BackgroundWorkTaskType) { guard !isDoingWorkInBackgound else { if expectedBackgroundWorkStatus?.rawValue ?? 0 < type.rawValue { log("Background work scheduled: \(type)") expectedBackgroundWorkStatus = type } return } DispatchQueue.global(qos: .utility).async { self.performAllBackgroundWorkItems(allItemsStartingAt: type) } } private func performAllBackgroundWorkItems(allItemsStartingAt type: BackgroundWorkTaskType) { didUpdateBackgroundItems = false expectedBackgroundWorkStatus = type log("Starting background task") while let type = setNextBackgroundWorkStatus() { log("Handling background task: \(type)") guard performAllItems(for: type) else { // If an error occurs, stop the background tasks backgroundTaskStatus = nil expectedBackgroundWorkStatus = nil break } } log("Background work completed") delegate?.databaseDidFinishBackgroundWork() } private func performAllItems(for type: BackgroundWorkTaskType) -> Bool { switch type { case .downloadCapNames: return downloadCapNames() case .downloadCounts: return downloadImageCounts() case .downloadClassifier: return downloadClassifier() case .uploadingCaps: return uploadCaps() case .uploadingImages: return uploadImages() case .downloadMainImages: return downloadMainImages() case .creatingThumbnails: return createThumbnails() case .creatingColors: return createColors() } } private func downloadCapNames() -> Bool { log("Downloading cap names") let result = DispatchGroup.singleTask { callback in download.names { names in guard let names = names else { callback(false) return } self.update(names: names) callback(true) } } log("Completed download of cap names") return result } private func downloadImageCounts() -> Bool { log("Downloading cap image counts") let result = DispatchGroup.singleTask { callback in download.imageCounts { counts in guard let counts = counts else { self.log("Failed to download server image counts") callback(false) return } let newCaps = self.didDownload(imageCounts: counts) guard newCaps.count > 0 else { callback(true) return } self.log("Found \(newCaps.count) new caps on the server.") self.downloadInfo(for: newCaps) { success in callback(success) } } } guard result else { log("Failed download of cap image counts") return false } log("Completed download of cap image counts") return true } private func downloadClassifier() -> Bool { log("Downloading classifier (if needed)") let result = DispatchGroup.singleTask { callback in download.classifierVersion { version in guard let version = version else { self.log("Failed to download server model version") callback(false) return } let ownVersion = self.classifierVersion guard ownVersion < version else { self.log("Not updating classifier: Own version \(ownVersion), server version \(version)") callback(true) return } let title = "Download classifier" let detail = ownVersion == 0 ? "A classifier to match caps is available for download (version \(version)). Would you like to download it now?" : "Version \(version) of the classifier is available for download (You have version \(ownVersion)). Would you like to download it now?" self.delegate!.database(needsUserConfirmation: title, body: detail) { proceed in guard proceed else { self.log("User skipped classifier download") callback(true) return } self.download.classifier { progress, received, total in let t = ByteCountFormatter.string(fromByteCount: total, countStyle: .file) let r = ByteCountFormatter.string(fromByteCount: received, countStyle: .file) let title = String(format: "%.0f", progress * 100) + " % (\(r) / \(t))" self.delegate?.database(completedBackgroundWorkItem: "Downloading classifier", subtitle: title) } completion: { url in guard let url = url else { self.log("Failed to download classifier") callback(false) return } let compiledUrl: URL do { compiledUrl = try MLModel.compileModel(at: url) } catch { self.log("Failed to compile downloaded classifier: \(error)") callback(false) return } guard self.storage.save(recognitionModelAt: compiledUrl) else { self.log("Failed to save compiled classifier") callback(false) return } callback(true) self.classifierVersion = version } } } } log("Downloaded classifier (if new version existed)") return result } private func uploadCaps() -> Bool { var completed = 0 while let cap = nextPendingCapUpload { guard upload.upload(cap) else { delegate?.database(didFailBackgroundWork: "Upload failed", subtitle: "Cap \(cap.id) not uploaded") return false } update(uploaded: true, for: cap.id) completed += 1 let total = completed + pendingCapUploadCount delegate?.database(completedBackgroundWorkItem: "Uploading caps", subtitle: "\(completed) of \(total)") } return true } private func uploadImages() -> Bool { var completed = 0 while let (id, version) = nextPendingImageUpload { guard let cap = self.cap(for: id) else { log("No cap \(id) to upload image \(version)") removePendingUpload(for: id, version: version) continue } guard let url = storage.existingImageUrl(for: cap.id, version: version) else { log("No image \(version) of cap \(id) to upload") removePendingUpload(for: id, version: version) continue } guard let count = upload.upload(imageAt: url, of: cap.id) else { delegate?.database(didFailBackgroundWork: "Upload failed", subtitle: "Image \(version) of cap \(id)") return false } if count > cap.count { update(count: count, for: cap.id) } removePendingUpload(for: id, version: version) completed += 1 let total = completed + pendingImageUploadCount delegate?.database(completedBackgroundWorkItem: "Uploading images", subtitle: "\(completed + 1) of \(total)") } return true } private func downloadMainImages() -> Bool { let missing = caps.map { $0.id }.filter { !storage.hasImage(for: $0) } let count = missing.count guard count > 0 else { log("No images to download") return true } log("Starting image downloads") let group = DispatchGroup() group.enter() var shouldDownload = true let title = "Download images" let detail = "\(count) caps have no image. Would you like to download them now? (~ \(ByteCountFormatter.string(fromByteCount: Int64(count * 10000), countStyle: .file))). Grid view is not available until all images are downloaded." delegate?.database(needsUserConfirmation: title, body: detail) { proceed in shouldDownload = proceed group.leave() } group.wait() guard shouldDownload else { log("User skipped image download") return false } group.enter() let queue = DispatchQueue(label: "images") let semaphore = DispatchSemaphore(value: 5) var downloadsAreSuccessful = true var completed = 0 for cap in missing { queue.async { guard downloadsAreSuccessful else { return } semaphore.wait() let url = self.storage.localImageUrl(for: cap) self.download.image(for: cap, to: url, queue: queue) { success in defer { semaphore.signal() } guard success else { self.delegate?.database(didFailBackgroundWork: "Download failed", subtitle: "Image of cap \(cap)") downloadsAreSuccessful = false group.leave() return } completed += 1 self.delegate?.database(completedBackgroundWorkItem: "Downloading images", subtitle: "\(completed) of \(missing.count)") if completed == missing.count { group.leave() } } } } guard group.wait(timeout: .now() + TimeInterval(missing.count * 2)) == .success else { log("Timed out downloading images") return false } log("Finished all image downloads") return true } private func createThumbnails() -> Bool { let missing = caps.map { $0.id }.filter { !storage.hasThumbnail(for: $0) } guard missing.count > 0 else { log("No thumbnails to create") return true } log("Creating thumbnails") let queue = DispatchQueue(label: "thumbnails") let semaphore = DispatchSemaphore(value: 5) let group = DispatchGroup() group.enter() var thumbnailsAreSuccessful = true var completed = 0 for cap in missing { queue.async { guard thumbnailsAreSuccessful else { return } semaphore.wait() defer { semaphore.signal() } guard let image = self.storage.image(for: cap) else { self.log("No image for cap \(cap) to create thumbnail") self.delegate?.database(didFailBackgroundWork: "Creation failed", subtitle: "Thumbnail of cap \(cap)") thumbnailsAreSuccessful = false group.leave() return } let thumb = Cap.thumbnail(for: image) guard self.storage.save(thumbnail: thumb, for: cap) else { self.log("Failed to save thumbnail for cap \(cap)") self.delegate?.database(didFailBackgroundWork: "Image not saved", subtitle: "Thumbnail of cap \(cap)") thumbnailsAreSuccessful = false group.leave() return } completed += 1 self.delegate?.database(completedBackgroundWorkItem: "Creating thumbnails", subtitle: "\(completed) of \(missing.count)") if completed == missing.count { group.leave() } } } guard group.wait(timeout: .now() + TimeInterval(missing.count * 2)) == .success else { log("Timed out creating thumbnails") return false } log("Finished all thumbnails") return true } private func createColors() -> Bool { let missing = capIds.subtracting(capsWithColors) guard missing.count > 0 else { log("No colors to create") return true } log("Creating colors") let queue = DispatchQueue(label: "colors") let semaphore = DispatchSemaphore(value: 5) let group = DispatchGroup() group.enter() var colorsAreSuccessful = true var completed = 0 for cap in missing { queue.async { guard colorsAreSuccessful else { return } semaphore.wait() defer { semaphore.signal() } guard let image = self.storage.ciImage(for: cap) else { self.log("No image for cap \(cap) to create color") self.delegate?.database(didFailBackgroundWork: "No thumbnail found", subtitle: "Color of cap \(cap)") colorsAreSuccessful = false group.leave() return } defer { self.context.clearCaches() } guard let color = image.averageColor(context: self.context) else { self.log("Failed to create color for cap \(cap)") self.delegate?.database(didFailBackgroundWork: "Calculation failed", subtitle: "Color of cap \(cap)") colorsAreSuccessful = false group.leave() return } guard self.set(color: color, for: cap) else { self.log("Failed to save color for cap \(cap)") self.delegate?.database(didFailBackgroundWork: "Color not saved", subtitle: "Color of cap \(cap)") colorsAreSuccessful = false group.leave() return } completed += 1 self.delegate?.database(completedBackgroundWorkItem: "Creating colors", subtitle: "\(completed) of \(missing.count)") if completed == missing.count { group.leave() } } } guard group.wait(timeout: .now() + TimeInterval(missing.count * 2)) == .success else { log("Timed out creating colors") return false } log("Finished all colors") return true } func hasNewClassifier(completion: @escaping (_ version: Int?, _ size: Int64?) -> Void) { download.classifierVersion { version in guard let version = version else { self.log("Failed to download server model version") completion(nil, nil) return } let ownVersion = self.classifierVersion guard ownVersion < version else { self.log("Not updating classifier: Own version \(ownVersion), server version \(version)") completion(nil, nil) return } self.log("Getting size of classifier \(version)") self.download.classifierSize { size in completion(version, size) } } } private func didDownload(imageCounts newCounts: [Int]) -> [Int : Int] { let capsCounts = capDict if newCounts.count != capsCounts.count { log("Downloaded \(newCounts.count) image counts, but \(capsCounts.count) caps stored locally") } var newCaps = [Int : Int]() let changed = newCounts.enumerated().compactMap { id, newCount -> Int? in let id = id + 1 guard let oldCount = capsCounts[id]?.count else { log("Received count \(newCount) for unknown cap \(id)") newCaps[id] = newCount return nil } guard oldCount != newCount else { return nil } self.update(count: newCount, for: id) return id } switch changed.count { case 0: log("Refreshed image counts for all caps without changes") case 1: log("Refreshed image counts for caps, changed cap \(changed[0])") case 2...10: log("Refreshed image counts for caps \(changed.map(String.init).joined(separator: ", ")).") default: log("Refreshed image counts for all caps (\(changed.count) changed)") } return newCaps } private func downloadInfo(for newCaps: [Int : Int], completion: @escaping (_ success: Bool) -> Void) { var success = true let group = DispatchGroup() for (id, count) in newCaps { group.enter() download.name(for: id) { name in guard let name = name else { self.log("Failed to get name for new cap \(id)") success = false group.leave() return } let cap = Cap(id: id, name: name, count: count) self.insert(cap: cap) group.leave() } } if group.wait(timeout: .now() + .seconds(30)) != .success { self.log("Timed out waiting for images to be downloaded") } completion(success) } func downloadImageCount(for cap: Int) { download.imageCount(for: cap) { count in guard let count = count else { return } self.update(count: count, for: cap) } } func setMainImage(of cap: Int, to version: Int) { guard version != 0 else { log("No need to switch main image with itself for cap \(cap)") return } upload.setMainImage(for: cap, to: version) { success in guard success else { self.log("Could not make \(version) the main image for cap \(cap)") return } guard self.storage.switchMainImage(to: version, for: cap) else { self.log("Could not switch \(version) to main image for cap \(cap)") return } DispatchQueue.main.async { self.delegate?.database(didLoadImageForCap: cap) } } } } extension Database: Logger { }