import Foundation import SwiftUI import Vision import CryptoKit final class Database: ObservableObject { private let imageCompressionQuality: CGFloat = 0.3 @AppStorage("classifier") private var storedLocalClassifierVersion = 0 { didSet { localClassifierVersion = storedLocalClassifierVersion } } @AppStorage("serverClassifier") private var storedServerClassifierVersion = 0 { didSet { serverClassifierVersion = storedServerClassifierVersion } } @Published private(set) var localClassifierVersion = 0 @Published private(set) var serverClassifierVersion = 0 @Published private(set) var classifierDownloadProgress: ClassifierProgress? var hasNewClassifier: Bool { serverClassifierVersion > localClassifierVersion } let images: ImageCache private let encoder = JSONEncoder() private let decoder = JSONDecoder() let serverUrl: URL let folderUrl: URL @AppStorage("authKey") private var serverAuthenticationKey: String = "" var hasServerAuthentication: Bool { serverAuthenticationKey != "" } @Published private(set) var caps: [Int : Cap] { didSet { scheduleSave() } } var nextCapId: Int { var next = 1 while caps[next] != nil { next += 1 } return next } @AppStorage("changed") private var changedCapStorage: String = "" private(set) var changedCaps: Set { get { Set(changedCapStorage.components(separatedBy: ",").compactMap(Int.init)) } set { changedCapStorage = newValue.map { "\($0)" }.joined(separator: ",") } } @AppStorage("uploads") private var pendingImageUploadStorage: String = "" private(set) var imageUploads: [Int: [Int]] { get { pendingImageUploadStorage.components(separatedBy: ";").reduce(into: [:]) { dict, string in let parts = string.components(separatedBy: "-") guard parts.count == 2 else { return } guard let cap = Int(parts[0]) else { return } dict[cap] = parts[1].components(separatedBy: ":").compactMap(Int.init) } } set { pendingImageUploadStorage = newValue.map { cap, images in "\(cap)-\(images.map { "\($0)" }.joined(separator: ":"))" }.joined(separator: ";") pendingImageUploadCount = newValue.values.reduce(0) { $0 + $1.count } } } @Published private(set) var pendingImageUploadCount = 0 private var uploadTimer: Timer? /// The classifications for all caps from the classifier @Published var matches = [Int : Float]() @Published var image: UIImage? = nil { didSet { classifyImage() } } private var classifier: Classifier? /** The time to wait for changes to be written to disk. This delay is used to prevent file writes for each small update to the caps. */ private let saveDelay: TimeInterval = 1 /** The time when a save should occur. No save is necessary if this property is `nil`. */ private var nextSaveTime: Date? @Published var isUploading = false @AppStorage("classifierClasses") private var _classifierClassesString: String = "" private var _classifierClassesCache: Set? private var classifierClasses: Set { get { _classifierClassesCache ?? loadClassifierClasses() } set { _classifierClassesCache = newValue DispatchQueue.main.async { self._classifierClassesString = newValue.map { "\($0)" }.joined(separator: ",") } } } private func loadClassifierClasses() -> Set { let elements: [Int] = _classifierClassesString.components(separatedBy: ",").compactMap { guard let id = Int($0) else { log("Failed to load classifier class from '\($0)'") return nil } return id } _classifierClassesCache = Set(elements) return _classifierClassesCache! } init(server: URL, folder: URL = FileManager.default.documentDirectory) { self.serverUrl = server self.folderUrl = folder self.caps = [:] let imageFolder = folder.appendingPathComponent("images") self.images = try! ImageCache( folder: imageFolder, server: server, thumbnailSize: CapsApp.thumbnailImageSize) self.localClassifierVersion = storedLocalClassifierVersion self.serverClassifierVersion = storedServerClassifierVersion ensureFolderExistence(gridStorageFolder) loadCaps() } func mainImage(for cap: Int) -> Int { caps[cap]?.mainImage ?? 0 } // MARK: URLs private var fm: FileManager { .default } private var localDbUrl: URL { folderUrl.appendingPathComponent("db.json") } private var localClassifierUrl: URL { folderUrl.appendingPathComponent("classifier.mlmodel") } private var imageUploadFolderUrl: URL { folderUrl.appendingPathComponent("uploads") } private var serverDbUrl: URL { serverUrl.appendingPathComponent("caps.json") } private var serverClassifierUrl: URL { serverUrl.appendingPathComponent("classifier.mlmodel") } private var serverClassifierClassesUrl: URL { serverUrl.appendingPathComponent("classifier.classes") } private var serverClassifierVersionUrl: URL { serverUrl.appendingPathComponent("version") } private var gridStorageFolder: URL { folderUrl.appendingPathComponent("grids") } func mainImageUrl(for cap: Int) -> URL? { guard let path = caps[cap]?.mainImagePath else { return nil } return serverUrl.appendingPathComponent(path) } // MARK: Disk storage private func loadCaps() { guard fm.fileExists(atPath: localDbUrl.path) else { return } let data: Data do { data = try Data(contentsOf: localDbUrl) } catch { log("Failed to read database file: \(error)") return } do { let array = try JSONDecoder().decode([Cap].self, from: data) self.caps = array.reduce(into: [:]) { $0[$1.id] = $1 } // Prevent immediate save after modifying caps nextSaveTime = nil } catch { log("Failed to decode database file: \(error)") return } } private func scheduleSave() { nextSaveTime = Date.now.addingTimeInterval(saveDelay) DispatchQueue.main.asyncAfter(deadline: .now() + saveDelay) { self.performScheduledSave() } } private func performScheduledSave() { guard let date = nextSaveTime else { // No save necessary, or already saved return } guard date < .now else { // Save pushed to future return } save() nextSaveTime = nil } private func save() { let data: Data do { data = try encoder.encode(caps.values.sorted()) } catch { log("Failed to encode database: \(error)") return } do { try data.write(to: localDbUrl) } catch { log("Failed to save database: \(error)") } log("Database saved") } @discardableResult private func ensureFolderExistence(_ url: URL) -> Bool { guard !fm.fileExists(atPath: url.path) else { return true } do { try fm.createDirectory(at: url, withIntermediateDirectories: true) return true } catch { log("Failed to create folder \(url.path): \(error)") return false } } // MARK: Downloads @discardableResult func downloadCaps() async -> Bool { log("Downloading cap data from \(serverDbUrl)") let data: Data let response: URLResponse do { (data, response) = try await URLSession.shared.data(from: serverDbUrl) } catch { log("Failed to download classifier version: \(error)") return false } guard let httpResponse = response as? HTTPURLResponse else { return false } guard httpResponse.statusCode == 200 else { log("Failed to download caps: \(httpResponse.statusCode)") return false } let capData: [CapData] do { capData = try decoder.decode([CapData].self, from: data) } catch { log("Failed to decode server database: \(error)") return false } var inserts = 0 var updates = 0 for cap in capData { guard var oldCap = caps[cap.id] else { caps[cap.id] = Cap(data: cap) inserts += 1 continue } guard oldCap != cap else { continue } if changedCaps.contains(oldCap.id) { #warning("Merge changed caps with server updates") } else { oldCap.update(with: cap) let save = oldCap DispatchQueue.main.async { self.caps[cap.id] = save } updates += 1 } } log("Updated database from server (\(inserts) added, \(updates) updated)") return true } @discardableResult func updateServerClassifierVersion() async -> Bool { let data: Data let response: URLResponse do { (data, response) = try await URLSession.shared.data(from: serverClassifierVersionUrl) } catch { log("Failed to download classifier version: \(error)") return false } guard (response as? HTTPURLResponse)?.statusCode == 200 else { return false } guard let string = String(data: data, encoding: .utf8) else { log("Classifier version is invalid data (not a string)") return false } guard let serverVersion = Int(string.trimmingCharacters(in: .whitespacesAndNewlines)) else { log("Classifier version has an invalid value '\(string)'") return false } DispatchQueue.main.async { self.storedServerClassifierVersion = serverVersion } return true } @discardableResult func downloadClassifier() async -> Bool { log("Downloading classifier") let progress = ClassifierProgress() DispatchQueue.main.async { self.classifierDownloadProgress = progress } defer { DispatchQueue.main.async { self.classifierDownloadProgress = nil } } let tempUrl: URL let response: URLResponse do { (tempUrl, response) = try await URLSession.shared.download(from: serverClassifierUrl, delegate: progress) } catch { log("Failed to download classifier version: \(error)") return false } guard (response as? HTTPURLResponse)?.statusCode == 200 else { return false } do { let url = self.localClassifierUrl if fm.fileExists(atPath: url.path) { try self.fm.removeItem(at: url) } try self.fm.moveItem(at: tempUrl, to: url) } catch { log("Failed to replace classifier: \(error)") return false } DispatchQueue.main.async { self.storedLocalClassifierVersion = self.serverClassifierVersion log("Downloaded classifier \(self.localClassifierVersion)") self.classifier = nil } return true } @discardableResult func downloadClassifierClasses() async -> Bool { log("Downloading classifier classes") let data: Data let response: URLResponse do { (data, response) = try await URLSession.shared.data(from: serverClassifierClassesUrl) } catch { log("Failed to download classifier classes: \(error)") return false } guard (response as? HTTPURLResponse)?.statusCode == 200 else { return false } guard let string = String(data: data, encoding: .utf8) else { log("Classifier classes is invalid data (not a string)") return false } let classes = string.components(separatedBy: ",") // Validate input var isValid = true let ids: [Int] = classes.compactMap { s in guard let id = Int(s) else { log("Invalid id '\(s)' in downloaded classes list") isValid = false return nil } if caps[id] == nil { // Caps which are deleted may still be recognized return nil } return id } guard isValid else { return false } self.classifierClasses = Set(ids) return true } /** Indicate that the cap has pending operations, such as determining the color or a thumbnail */ func hasPendingOperations(for cap: Int) -> Bool { return false } func cap(for id: Int) -> Cap? { caps[id] } // MARK: Adding new data func save(newCap name: String) -> Cap { let cap = Cap(id: nextCapId, name: name) caps[cap.id] = cap DispatchQueue.main.async { self.changedCaps.insert(cap.id) } return cap } @discardableResult func save(_ image: UIImage, for capId: Int) -> Bool { guard let cap = caps[capId] else { log("Failed to save image for missing cap \(capId)") return false } guard images.save(image, for: CapImage(cap: cap.id, version: cap.imageCount)) else { return false } log("Saved image \(cap.imageCount) for cap \(capId)") if imageUploads[capId] != nil { DispatchQueue.main.async { self.imageUploads[capId]!.append(cap.imageCount) } } else { DispatchQueue.main.async { self.imageUploads[capId] = [cap.imageCount] } } DispatchQueue.main.async { self.caps[capId]!.imageCount += 1 } return true } func update(name: String, for capId: Int) -> Bool { guard var cap = caps[capId] else { log("Failed to update name for missing cap \(capId)") return false } cap.name = name caps[capId] = cap changedCaps.insert(capId) return true } // MARK: Uploads func startRegularUploads() { guard uploadTimer == nil else { return } log("Starting upload timer") DispatchQueue.main.async { self.uploadTimer = Timer.scheduledTimer(withTimeInterval: 5, repeats: true, block: self.uploadTimerElapsed) } } private func uploadTimerElapsed(timer: Timer) { Task { await uploadAll() } } private func uploadAll() async { guard !isUploading else { log("Already uploading") return } DispatchQueue.main.async { self.isUploading = true } defer { DispatchQueue.main.async { self.isUploading = false } } guard !changedCaps.isEmpty || pendingImageUploadCount > 0 else { return } log("Starting uploads") let uploaded = await uploadAllChangedCaps() DispatchQueue.main.async { self.changedCaps.subtract(uploaded) } await uploadAllImages() log("Uploads finished") } /** Indicate that the cap has pending uploads, either changes or images */ func hasPendingUpdates(for cap: Int) -> Bool { changedCaps.contains(cap) || imageUploads[cap] != nil } private func capId(from url: URL) -> Int? { Int(url.lastPathComponent.components(separatedBy: "-").first!) } private func uploadAllImages() async { guard hasServerAuthentication else { log("No server authentication to upload to server") return } for (cap, images) in imageUploads { for image in images { guard let url = self.images.availableImageUrl(CapImage(cap: cap, version: image)) else { log("Missing upload image \(image) for cap \(cap)") continue } guard await upload(imageAt: url, for: cap) else { log("Failed to upload image \(url.lastPathComponent)") continue } log("Uploaded image \(image) for cap \(cap)") let remaining = imageUploads[cap]?.filter { $0 != image } if let r = remaining, !r.isEmpty { DispatchQueue.main.async { self.imageUploads[cap] = r } } else { DispatchQueue.main.async { self.imageUploads[cap] = nil } } } } } @discardableResult private func upload(imageAt url: URL, for cap: Int) async -> Bool { guard hasServerAuthentication else { return false } guard let data = try? Data(contentsOf: url) else { log("No image data found for image \(url.lastPathComponent) (Cap \(cap))") return false } let url = serverUrl .appendingPathComponent("image") .appendingPathComponent("\(cap)") var request = URLRequest(url: url) request.addValue(serverAuthenticationKey, forHTTPHeaderField: "key") request.httpMethod = "POST" do { let (_, response) = try await URLSession.shared.upload(for: request, from: data) guard let httpResponse = response as? HTTPURLResponse else { log("Unexpected response for upload of image \(url.lastPathComponent): \(response)") return false } guard httpResponse.statusCode == 200 else { log("Failed to upload image \(url.path): Response \(httpResponse.statusCode)") return false } return true } catch { log("Failed to upload image \(url.lastPathComponent): \(error)") return false } } var pendingCapUploadCount: Int { changedCaps.count } private func uploadAllChangedCaps() async -> Set { guard hasServerAuthentication else { log("No server authentication to upload to server") return .init() } var uploaded = Set() for capId in changedCaps { guard let cap = caps[capId] else { log("Missing cap \(capId) to upload") uploaded.insert(capId) continue } guard await upload(cap: cap) else { continue } log("Uploaded cap \(capId)") uploaded.insert(capId) } return uploaded } @discardableResult private func upload(cap: Cap) async -> Bool { guard hasServerAuthentication else { return false } let data: Data do { /// `Cap` and `CapData` have equivalent JSON layout data = try encoder.encode(cap.data) } catch { log("Failed to encode cap \(cap.id) for upload: \(error)") return false } let url = serverUrl .appendingPathComponent("cap") var request = URLRequest(url: url) request.httpMethod = "POST" request.addValue(serverAuthenticationKey, forHTTPHeaderField: "key") do { let (_, response) = try await URLSession.shared.upload(for: request, from: data) guard let httpResponse = response as? HTTPURLResponse else { log("Unexpected response for upload of cap \(cap.id): \(response)") return false } guard httpResponse.statusCode == 200 else { log("Failed to upload cap \(cap.id): Response \(httpResponse.statusCode)") return false } DispatchQueue.main.async { self.changedCaps.remove(cap.id) } return true } catch { log("Failed to upload cap \(cap.id): \(error)") return false } } func setMainImage(_ version: Int, for capId: Int) async -> Cap? { guard hasServerAuthentication else { log("No authorization to set main image") return nil } guard var cap = cap(for: capId) else { log("No cap \(capId) to set main image") return nil } guard version < cap.imageCount else { log("Invalid main image \(version) for \(capId) with only \(cap.imageCount) images") return nil } cap.mainImage = version let finalCap = cap DispatchQueue.main.async { self.caps[capId] = finalCap log("Set main image \(version) for \(capId)") } return finalCap } func delete(image: CapImage) async -> Bool { guard hasServerAuthentication else { log("No authorization to delete cap image") return false } guard let cap = cap(for: image.cap) else { log("No cap \(image.cap) to delete cap image") return false } guard image.version < cap.imageCount else { log("Invalid image \(image.version) to delete for \(cap.id) with only \(cap.imageCount) images") return false } let url = serverUrl .appendingPathComponent("delete/\(cap.id)/\(image.version)") var request = URLRequest(url: url) request.httpMethod = "POST" request.addValue(serverAuthenticationKey, forHTTPHeaderField: "key") do { let (data, response) = try await URLSession.shared.data(for: request) guard let httpResponse = response as? HTTPURLResponse else { log("Unexpected response deleting image \(image.version) of cap \(cap.id): \(response)") return false } guard httpResponse.statusCode == 200 else { log("Failed to delete image \(image.version) of cap \(cap.id): Response \(httpResponse.statusCode)") return false } let newCap: Cap do { newCap = try JSONDecoder().decode(Cap.self, from: data) } catch { log("Invalid response data deleting image \(image.version) of cap \(cap.id): \(data)") return false } // Delete cached images images.removeCachedImages(for: cap.id) // Update cap caps[newCap.id] = newCap log("Deleted image \(image.version) of cap \(cap.id)") return true } catch { log("Failed to delete image \(image.version) of cap \(cap.id): \(error)") return false } } func delete(cap: Int) async -> Bool { guard hasServerAuthentication else { log("No authorization to delete cap") return false } guard caps[cap] != nil else { log("No cap \(cap) to delete") return false } let url = serverUrl.appendingPathComponent("delete/\(cap)") var request = URLRequest(url: url) request.httpMethod = "POST" request.addValue(serverAuthenticationKey, forHTTPHeaderField: "key") do { let (_, response) = try await URLSession.shared.data(for: request) guard let httpResponse = response as? HTTPURLResponse else { log("Unexpected response deleting cap \(cap): \(response)") return false } guard httpResponse.statusCode == 200 else { log("Failed to delete cap \(cap): Response \(httpResponse.statusCode)") return false } // Delete cached images images.removeCachedImages(for: cap) // Delete cap DispatchQueue.main.async { self.caps[cap] = nil } log("Deleted cap \(cap)") return true } catch { log("Failed to delete cap \(cap): \(error)") return false } } // MARK: Classification /// The compiled recognition model on disk private var recognitionModel: VNCoreMLModel? { guard fm.fileExists(atPath: localClassifierUrl.path) else { log("No recognition model to load from disk") return nil } do { log("Loading model from disk") let newUrl = try MLModel.compileModel(at: localClassifierUrl) let model = try MLModel(contentsOf: newUrl) return try VNCoreMLModel(for: model) } catch { log("Failed to load recognition model: \(error)") return nil } } private func classifyImage() { guard let image = image?.cgImage else { matches.removeAll() log("Image removed") return } DispatchQueue.global().async { guard let classifier = self.getClassifier() else { return } log("Image classification started") classifier.recognize(image: image) { matches in DispatchQueue.main.async { self.matches = matches ?? [:] } } } } func canClassify(cap id: Int) -> Bool { classifierClasses.contains(id) } private func getClassifier() -> Classifier? { if let classifier = classifier { return classifier } guard let model = recognitionModel else { return nil } return Classifier(model: model) } // MARK: Grid var availableGrids: [String] { do { return try fm.contentsOfDirectory(at: gridStorageFolder, includingPropertiesForKeys: nil) .filter { $0.pathExtension == "caps" } .map { $0.deletingPathExtension().lastPathComponent } } catch { log("Failed to load available grids: \(error)") return [] } } private func gridFileUrl(_ grid: String) -> URL { gridStorageFolder.appendingPathComponent(grid).appendingPathExtension("caps") } private func gridImageUrl(_ grid: String) -> URL { gridStorageFolder.appendingPathComponent(grid).appendingPathExtension("jpg") } func load(grid: String) -> ImageGrid? { let url = gridFileUrl(grid) guard fm.fileExists(atPath: url.path) else { return nil } do { let data = try Data(contentsOf: url) var loaded = try decoder.decode(ImageGrid.self, from: data) // Add all missing caps to the end of the image let newCaps = Set(caps.keys).subtracting(loaded.capPlacements).sorted() loaded.capPlacements += newCaps log("Grid \(grid) loaded (\(newCaps.count) new caps)") return loaded } catch { log("Failed to load grid \(grid): \(error)") return nil } } @discardableResult func save(_ grid: ImageGrid, named name: String) -> Bool { let url = gridFileUrl(name) do { let data = try encoder.encode(grid) try data.write(to: url) log("Grid \(name) saved") return true } catch { log("Failed to save grid \(name): \(error)") return false } } // MARK: Grid images func load(gridImage: String) -> UIImage? { let url = gridImageUrl(gridImage) return UIImage(at: url) } @discardableResult func save(gridImage: UIImage, for grid: String) -> Bool { guard let data = gridImage.jpegData(compressionQuality: 0.9) else { return false } let url = gridImageUrl(grid) do { try data.write(to: url) return true } catch { log("Failed to save grid image \(grid): \(error)") return false } } // MARK: Statistics var numberOfCaps: Int { caps.count } var numberOfImages: Int { caps.values.reduce(0) { $0 + $1.imageCount } } var averageImageCount: Float { Float(numberOfImages) / Float(numberOfCaps) } var classifierClassCount: Int { classifierClasses.count } func imageCacheSize() -> Int { fm.directorySize(images.folder) } var databaseSize: Int { localDbUrl.fileSize } var classifierSize: Int { localClassifierUrl.fileSize } func clearImageCache() { let imagesToKeep = caps.values.map { $0.image } images.clearImageCache(keeping: Set(imagesToKeep)) } } extension Database { final class ClassifierProgress: NSObject, ObservableObject { @Published var bytesLoaded: Double @Published var total: Double var percentage: Double { guard total > 0 else { return 0.0 } return bytesLoaded * 100 / total } init(bytesLoaded: Double = 0, total: Double = 0) { self.bytesLoaded = bytesLoaded self.total = total } } } extension Database.ClassifierProgress: URLSessionDownloadDelegate { func urlSession(_ session: URLSession, downloadTask: URLSessionDownloadTask, didFinishDownloadingTo location: URL) { } func urlSession(_ session: URLSession, downloadTask: URLSessionDownloadTask, didWriteData bytesWritten: Int64, totalBytesWritten: Int64, totalBytesExpectedToWrite: Int64) { DispatchQueue.main.async { self.bytesLoaded = Double(totalBytesWritten) self.total = Double(totalBytesExpectedToWrite) } } } extension Database { static var mock: Database { let db = Database(server: URL(string: "https://christophhagen.de/caps")!) db.caps = [ Cap(id: 123, name: "My new cap"), Cap(id: 234, name: "My favorite cap"), Cap(id: 345, name: "My oldest cap"), Cap(id: 456, name: "My new cap"), Cap(id: 567, name: "My favorite cap"), Cap(id: 678, name: "My oldest cap"), ].reduce(into: [:]) { $0[$1.id] = $1 } db.image = UIImage(systemSymbol: .photo) return db } static var largeMock: Database { let db = Database(server: URL(string: "https://christophhagen.de/caps")!) db.caps = (1..<500) .map { Cap(id: $0, name: "Cap \($0)") } .reduce(into: [:]) { $0[$1.id] = $1 } db.image = UIImage(systemSymbol: .photo) return db } }