import Foundation
import SwiftUI
import Vision
import CryptoKit

final class Database: ObservableObject {

    @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()

    @AppStorage("serverUrl")
    var serverPath: String = "" {
        didSet {
            images.server = serverUrl
        }
    }

    var serverUrl: URL? {
        .init(string: serverPath)
    }

    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<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: ";")
            updatePendingImageUploadCount(imageUploads: newValue)
        }
    }

    private func updatePendingImageUploadCount(imageUploads: [Int: [Int]]) {
        pendingImageUploadCount = imageUploads.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<Int>?
    
    private var classifierClasses: Set<Int> {
        get {
            _classifierClassesCache ?? loadClassifierClasses()
        }
        set {
            _classifierClassesCache = newValue
            DispatchQueue.main.async {
                self._classifierClassesString = newValue.map { "\($0)" }.joined(separator: ",")
            }
        }
    }
    
    private func loadClassifierClasses() -> Set<Int> {
        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(folder: URL = FileManager.default.documentDirectory) {
        self.folderUrl = folder
        self.caps = [:]

        let imageFolder = folder.appendingPathComponent("images")
        self.images = try! ImageCache(
            folder: imageFolder,
            server: nil,
            thumbnailSize: CapsApp.thumbnailImageSize)
        
        self.localClassifierVersion = storedLocalClassifierVersion
        self.serverClassifierVersion = storedServerClassifierVersion

        images.server = serverUrl
        ensureFolderExistence(gridStorageFolder)
        loadCaps()
        updatePendingImageUploadCount(imageUploads: imageUploads)
    }

    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 {
        guard let serverDbUrl else {
            log("No server url set to download cap data")
            return false
        }
        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 {
        guard let serverClassifierVersionUrl else {
            log("No server url to download classifier version")
            return false
        }
        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 {
        guard let serverClassifierUrl else {
            log("No server url to download classifier")
            return false
        }
        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 {
        guard let serverClassifierClassesUrl else {
            log("No server url to download classifier classes")
            return false
        }
        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 let serverUrl else {
            log("No server url to upload image")
            return false
        }
        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 serverImageUrl = serverUrl
            .appendingPathComponent("image")
            .appendingPathComponent("\(cap)")
        var request = URLRequest(url: serverImageUrl)
        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 {
                if httpResponse.statusCode == 410 {
                    log("Missing cap for image \(url.lastPathComponent), reupload cap")
                    // Missing cap, force upload
                    DispatchQueue.main.async {
                        self.changedCaps.insert(cap)
                    }
                } 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 let serverUrl else {
            log("No server url to upload cap")
            return false
        }
        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 let serverUrl else {
            log("No server url to delete image")
            return false
        }
        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 let serverUrl else {
            log("No server url to delete cap")
            return false
        }
        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()
        db.serverPath = "https://caps.christophhagen.de"
        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()
        db.serverPath = "https://caps.christophhagen.de"
        db.caps = (1..<500)
            .map { Cap(id: $0, name: "Cap \($0)") }
            .reduce(into: [:]) { $0[$1.id] = $1 }
        db.image = UIImage(systemSymbol: .photo)
        return db
    }
}