Caps-iOS/Caps/Data/Database.swift

import Foundation
import SwiftUI
import Vision
import CryptoKit

final class Database: ObservableObject {

    private let imageCompressionQuality: CGFloat = 0.3

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

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

    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 {
        (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: ",")
        }
    }
    
    @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: ";")
        }
    }

    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

    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)

        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 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 {
            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")
    }

    @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 {
        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)
                let save = oldCap
                DispatchQueue.main.async {
                    self.caps[cap.id] = save
                }
                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 (Local: \(classifierVersion) Server: \(serverVersion))")
            return false
        }
        print("New classifier available (Local: \(classifierVersion) Server: \(serverVersion))")
        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: nil)
        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
    }

    var pendingImageUploadCount: Int {
        imageUploads.values.reduce(0) { $0 + $1.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 (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 {
            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: Grid

    var availableGrids: [String] {
        do {
            return try fm.contentsOfDirectory(at: gridStorageFolder, includingPropertiesForKeys: nil)
                .filter { $0.pathExtension == "caps" }
                .map { $0.deletingPathExtension().lastPathComponent }
        } catch {
            print("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
            print("Grid \(grid) loaded (\(newCaps.count) new caps)")
            return loaded
        } catch {
            print("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)
            print("Grid \(name) saved")
            return true
        } catch {
            print("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 {
            print("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 {
        let version = classifierVersion
        return caps.values.filter { $0.classifiable(by: version) }.count
    }

    func imageCacheSize() async -> Int {
        fm.directorySize(images.folder)
    }

    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
    }

    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)", classifier: nil)}
            .reduce(into: [:]) { $0[$1.id] = $1 }
        db.image = UIImage(systemSymbol: .photo)
        return db
    }
}