Caps-iOS/Caps/Data/Database.swift

import Foundation
import SwiftUI
import Vision
import CryptoKit

final class Database: ObservableObject {

    static let imageCacheMemory = 10_000_000

    static let imageCacheStorage = 200_000_000

    private let imageCompressionQuality: CGFloat = 0.3

    private static var documentDirectory: URL {
        try! FileManager.default.url(
        for: .documentDirectory,
        in: .userDomainMask,
        appropriateFor: nil, create: true)
    }

    private var fm: FileManager {
        .default
    }

    private var localDbUrl: URL {
        Database.documentDirectory.appendingPathComponent("db.json")
    }

    private var localClassifierUrl: URL {
        Database.documentDirectory.appendingPathComponent("classifier.mlmodel")
    }

    private var imageUploadFolderUrl: URL {
        Database.documentDirectory.appendingPathComponent("uploads")
    }

    private var serverDbUrl: URL {
        serverUrl.appendingPathComponent("caps.json")
    }

    private var serverClassifierUrl: URL {
        serverUrl.appendingPathComponent("classifier.mlmodel")
    }

    private var serverClassifierVersionUrl: URL {
        serverUrl.appendingPathComponent("classifier.version")
    }

    private let encoder = JSONEncoder()
    private let decoder = JSONDecoder()

    let serverUrl: URL

    @AppStorage("authKey")
    private var serverAuthenticationKey: String = ""

    var hasServerAuthentication: Bool {
        serverAuthenticationKey != ""
    }

    @Published
    private(set) var caps: [Int : Cap] {
        didSet { scheduleSave() }
    }

    var nextCapId: Int {
        (caps.values.max()?.id ?? 0) + 1
    }

    @AppStorage("changed")
    private var changedCapStorage: String = ""

    private(set) var changedCaps: Set<Int> {
        get {
            Set(changedCapStorage.components(separatedBy: ",").compactMap(Int.init))
        }
        set {
            changedCapStorage = newValue.map { "\($0)" }.joined(separator: ",")
        }
    }

    private lazy var imageUploads: [Int: Int] = loadImageUploadCounts()

    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?

    let imageCache: URLCache

    init(server: URL) {
        self.serverUrl = server
        self.caps = [:]

        let cacheDirectory = Database.documentDirectory.appendingPathComponent("images")
        self.imageCache = URLCache(
            memoryCapacity: Database.imageCacheMemory,
            diskCapacity: Database.imageCacheStorage,
            directory: cacheDirectory)
        loadCaps()
    }

    @Published
    var isUploading = false

    // MARK: Disk storage

    private func loadCaps() {
        guard fm.fileExists(atPath: localDbUrl.path) else {
            return
        }
        let data: Data
        do {
            data = try Data(contentsOf: localDbUrl)
        } catch {
            print("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 {
            print("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 {
            print("Failed to encode database: \(error)")
            return
        }
        do {
            try data.write(to: localDbUrl)
        } catch {
            print("Failed to save database: \(error)")
        }
        print("Database saved")
    }

    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 {
        print("Downloading cap data")
        let data: Data
        let response: URLResponse
        do {
            (data, response) = try await URLSession.shared.data(from: serverDbUrl)
        } catch {
            print("Failed to download classifier version: \(error)")
            return false
        }
        guard (response as? HTTPURLResponse)?.statusCode == 200 else {
            return false
        }

        let capData: [CapData]
        do {
            capData = try decoder.decode([CapData].self, from: data)
        } catch {
            print("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)
                caps[cap.id] = oldCap
                updates += 1
            }
        }
        print("Updated database from server (\(inserts) added, \(updates) updated)")
        return true
    }

    @discardableResult
    func serverHasNewClassifier() async -> Bool {
        let data: Data
        let response: URLResponse
        do {
            (data, response) = try await URLSession.shared.data(from: serverClassifierVersionUrl)
        } catch {
            print("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.serverClassifierVersion = serverVersion
        }
        guard serverVersion > self.classifierVersion else {
            print("No new classifier available")
            return false
        }
        print("New classifier \(serverVersion) available")
        return true
    }

    @discardableResult
    func downloadClassifier() async -> Bool {
        print("Downloading classifier")
        let tempUrl: URL
        let response: URLResponse
        do {
            (tempUrl, response) = try await URLSession.shared.download(from: serverClassifierUrl)
        } catch {
            print("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 {
            print("Failed to replace classifier: \(error)")
            return false
        }
        DispatchQueue.main.async {
            self.classifierVersion = self.serverClassifierVersion
            self.classifier = nil
        }
        print("Downloaded classifier \(classifierVersion)")
        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
    }

    // MARK: Adding new data

    func save(newCap name: String) -> Cap {
        let cap = Cap(id: nextCapId, name: name, classifier: serverClassifierVersion)
        caps[cap.id] = cap
        DispatchQueue.main.async {
            self.changedCaps.insert(cap.id)
        }
        return cap
    }

    @discardableResult
    func save(_ image: UIImage, for capId: Int) -> Bool {
        guard caps[capId] != nil else {
            log("Failed to save image for missing cap \(capId)")
            return false
        }
        guard ensureFolderExistence(imageUploadFolderUrl) else {
            return false
        }
        guard let data = image.jpegData(compressionQuality: imageCompressionQuality) else {
            log("Failed to compress image for cap: \(capId)")
            return false
        }
        let hash = Data(SHA256.hash(data: data)).hexEncoded.prefix(16)
        let url = imageUploadFolderUrl.appendingPathComponent("\(capId)-\(hash).jpg")
        do {
            try data.write(to: url)
        } catch {
            log("Failed to save \(url.lastPathComponent): \(error)")
            return false
        }
        log("Saved \(url.lastPathComponent) for upload")
        caps[capId]?.imageCount += 1
        return true
    }

    private func loadImageUploadCounts() -> [Int : Int] {
        var result = [Int : Int]()
        pendingImageUploads.forEach { url in
            guard let capId = capId(from: url) else {
                return
            }
            if let old = result[capId] {
                result[capId] = old + 1
            } else {
                result[capId] = 1
            }
        }
        return result
    }

    // 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
        }
        log("Starting uploads")
        let uploaded = await uploadAllChangedCaps()
        DispatchQueue.main.async {
            self.changedCaps.subtract(uploaded)
        }
        await uploadAllImages()
        log("Uploads finished")
        DispatchQueue.main.async {
            self.isUploading = false
        }
    }

    /**
     Indicate that the cap has pending uploads, either changes or images
     */
    func hasPendingUpdates(for cap: Int) -> Bool {
        changedCaps.contains(cap) || imageUploads[cap] != nil
    }

    private var pendingImageUploads: [URL] {
        (try? fm.contentsOfDirectory(at: imageUploadFolderUrl, includingPropertiesForKeys: nil)) ?? []
    }

    var pendingImageUploadCount: Int {
        pendingImageUploads.count
    }

    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 url in pendingImageUploads {
            guard let capId = capId(from: url) else {
                log("Unexpected image \(url.lastPathComponent) in upload folder")
                continue
            }
            guard fm.fileExists(atPath: url.path) else {
                log("Missing image \(url.lastPathComponent) in upload folder")
                continue
            }
            guard await upload(imageAt: url, for: capId) else {
                log("Failed to upload image \(url.lastPathComponent)")
                continue
            }
            log("Uploaded image \(url.lastPathComponent)")
            do {
                try fm.removeItem(at: url)
            } catch {
                log("Failed to remove uploaded image \(url.lastPathComponent): \(error)")
            }
        }
    }

    @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 {
            return false
        }
        let url = serverUrl
            .appendingPathComponent("images")
            .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.lastPathComponent): 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<Int> {
        guard hasServerAuthentication else {
            log("No server authentication to upload to server")
            return .init()
        }
        var uploaded = Set<Int>()
        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
        }
    }

    // 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 ?? [:]
                }
            }
        }
    }

    private func getClassifier() -> Classifier? {
        if let classifier = classifier {
            return classifier
        }
        guard let model = recognitionModel else {
            return nil
        }
        return Classifier(model: model)
    }

    // MARK: Statistics

    var numberOfCaps: Int {
        caps.count
    }

    var numberOfImages: Int {
        caps.values.reduce(0) { $0 + $1.imageCount }
    }

    var averageImageCount: Float {
        Float(numberOfImages) / Float(numberOfCaps)
    }

    @AppStorage("classifier")
    private(set) var classifierVersion = 0

    @AppStorage("serverClassifier")
    private(set) var serverClassifierVersion = 0

    var classifierClassCount: Int {
        let version = classifierVersion
        return caps.values.filter { $0.classifiable(by: version) }.count
    }

    var imageCacheSize: Int {
        imageCache.currentDiskUsage
    }

    var databaseSize: Int {
        localDbUrl.fileSize
    }

    var classifierSize: Int {
        localClassifierUrl.fileSize
    }
}

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
    }
}