Caps-Server/Sources/App/CapServer.swift

import Foundation
import Vapor
import Clairvoyant

final class CapServer: ServerOwner {
    
    // MARK: Paths
    
    private let imageFolder: URL

    /// The file where the database of caps is stored
    private let dbFile: URL

    /// The file to store the HTML info of the cap count
    private let htmlFile: URL
    
    private let classifierVersionFile: URL
    
    private let classifierFile: URL
    
    private let fm = FileManager.default
    
    // MARK: Caps

    private var writers: Set<String>
    
    var classifierVersion: Int = 0 {
        didSet {
            writeClassifierVersion()
            updateMonitoredClassifierVersionProperty()
        }
    }

    /**
     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?

    private var caps = [Int: Cap]() {
        didSet {
            scheduleSave()
            updateMonitoredPropertiesOnCapChange()
        }
    }

    var nextClassifierVersion: Int {
        caps.values.compactMap { $0.classifierVersion }.max() ?? 1
    }

    var capCount: Int {
        caps.count
    }

    var imageCount: Int {
        caps.reduce(0) { $0 + $1.value.count }
    }

    init(in folder: URL, writers: [String]) {
        self.imageFolder = folder.appendingPathComponent("images")
        self.dbFile = folder.appendingPathComponent("caps.json")
        self.htmlFile = folder.appendingPathComponent("count.html")
        self.classifierVersionFile = folder.appendingPathComponent("classifier.version")
        self.classifierFile = folder.appendingPathComponent("classifier.mlmodel")
        self.writers = Set(writers)
    }

    func loadData() throws {
        loadClassifierVersion(at: classifierVersionFile)
        try loadCaps()
        try updateCounts()
        saveCapCountHTML()
    }

    private func loadClassifierVersion(at url: URL) {
        guard fm.fileExists(atPath: url.path) else {
            return
        }
        let content: String
        do {
            content = try String(contentsOf: url)
                .trimmingCharacters(in: .whitespacesAndNewlines)
        } catch {
            log("Failed to read classifier version file: \(error)")
            return
        }

        guard let value = Int(content) else {
            log("Invalid classifier version: \(content)")
            return
        }
        self.classifierVersion = value
    }

    private func writeClassifierVersion() {
        do {
            try "\(classifierVersion)".data(using: .utf8)!
                .write(to: classifierVersionFile)
        } catch {
            log("Failed to save classifier version: \(error)")
        }
    }

    private func loadCaps() throws {
        do {
            let data = try Data(contentsOf: dbFile)
            caps = try JSONDecoder().decode([Cap].self, from: data)
                .reduce(into: [:]) { $0[$1.id] = $1 }
        } catch {
            log("Failed to load caps: \(error)")
            throw error
        }
        log("\(caps.count) caps loaded")
    }

    private func scheduleSave() {
        nextSaveTime = Date().addingTimeInterval(saveDelay)
        DispatchQueue.global().asyncAfter(deadline: .now() + saveDelay) {
            self.performScheduledSave()
        }
    }

    private func performScheduledSave() {
        guard let date = nextSaveTime else {
            // No save necessary, or already saved
            return
        }
        guard date < Date() else {
            // Save pushed to future
            return
        }
        do {
            try saveCaps()
            nextSaveTime = nil
        } catch {
            // Attempt save again
            scheduleSave()
        }
    }

    private func saveCaps() throws {
        let data = try JSONEncoder().encode(caps.values.sorted())
        try data.write(to: dbFile)
    }

    private func saveCapCountHTML() {
        let count = caps.count
        let content =
        """
        <body style="margin: 0;">
            <div style="display: flex; justify-content: center;">
                <div style="font-size: 60px; font-family: 'SF Pro Display',-apple-system,BlinkMacSystemFont,Helvetica,sans-serif; -webkit-font-smoothing: antialiased;">\(count)</div>
            </div>
        </body>
        """
        try? content.data(using: .utf8)!.write(to: htmlFile)
    }

    // MARK: Paths

    func folder(of cap: Int) -> URL {
        imageFolder.appendingPathComponent(String(format: "%04d", cap))
    }

    func file(of cap: Int, version: Int) -> URL {
        folder(of: cap).appendingPathComponent(String(format: "%04d-%02d.jpg", cap, version))
    }

    // MARK: Authentication

    func hasAuthorization(for key: String) -> Bool {
        // Note: This is not a constant-time compare, so there may be an opportunity
        // for timing attack here. Sets perform hashed lookups, so this may be less of an issue,
        // and we're not doing anything critical in this application.
        // Worst case, an unauthorized person with a lot of free time and energy to hack this system
        // is able to change contents of the database, which are backed up in any case.
        writers.contains(key)
    }
    
    // MARK: Counts

    private func updateCounts() throws {
        do {
            caps = try caps.mapValues {
                var cap = $0
                cap.count = try count(of: $0.id)
                return cap
            }
        } catch {
            log("Failed to update counts: \(error)")
            throw error
        }
    }

    func countImages(in folder: URL) throws -> Int {
        try fm.contentsOfDirectory(at: folder, includingPropertiesForKeys: nil)
            .filter({ $0.pathExtension == "jpg" }).count
    }

    /**
     Get the image count of a cap.
     */
    func count(of cap: Int) throws -> Int {
        let f = folder(of: cap)
        guard fm.fileExists(atPath: f.path) else {
            return 0
        }
        return try countImages(in: f)
    }

    // MARK: Images
    
    /**
     Save a cap image to disk.
     
     Automatically creates the image name with the current image count.
     - Parameter data: The image data
     - Parameter cap: The id of the cap.
     - Throws: `CapError.unknownId`, if the cap doesn't exist. `CapError.dataInconsistency` if an image already exists for the current count.
     */
    func save(image data: Data, for cap: Int) throws {
        guard caps[cap] != nil else {
            throw CapError.unknownId
        }
        var id = 0
        let capFolder = folder(of: cap)
        var f = file(of: cap, version: id)
        if fm.fileExists(atPath: capFolder.path) {
            while fm.fileExists(atPath: f.path) {
                id += 1
                f = file(of: cap, version: id)
            }
        } else {
            try fm.createDirectory(at: capFolder, withIntermediateDirectories: true)
        }
        try data.write(to: f)
        caps[cap]!.count = try count(of: cap)
        log("Added image \(id) for cap \(cap)")
    }
    
    func switchMainImage(to version: Int, for cap: Int) throws {
        let file2 = file(of: cap, version: version)
        guard fm.fileExists(atPath: file2.path) else {
            log("No image \(version) for cap \(cap)")
            throw CapError.invalidFile
        }
        caps[cap]?.mainImage = version
        log("Switched cap \(cap) to version \(version)")
    }

    func addOrUpdate(_ cap: Cap) throws {
        if let existingCap = caps[cap.id] {
            update(existingCap, with: cap)
        } else {
            try add(cap)
        }
    }

    private func add(_ cap: Cap) throws {
        guard cap.mainImage == 0 else {
            throw CapError.invalidData
        }
        var cap = cap
        cap.count = 0
        cap.classifierVersion = nextClassifierVersion
        caps[cap.id] = cap
        saveCapCountHTML()
        log("Added cap \(cap.id) '\(cap.name)'")
    }

    private func update(_ existingCap: Cap, with cap: Cap) {
        var updatedCap = existingCap
        if cap.name != "" {
            updatedCap.name = cap.name
        }
        let url = file(of: existingCap.id, version: cap.mainImage)
        if fm.fileExists(atPath: url.path) {
            updatedCap.mainImage = cap.mainImage
        }
        if let color = cap.color {
            updatedCap.color = color
        }
        caps[existingCap.id] = updatedCap
        log("Updated cap \(existingCap.id)")
    }
    
    func updateTrainedClasses(content: String) {
        let trainedCaps = content
            .components(separatedBy: "\n")
            .compactMap(Int.init)
        let version = classifierVersion
        for cap in trainedCaps {
            if caps[cap]?.classifierVersion == nil {
                caps[cap]?.classifierVersion = version
            }
        }
        log("Updated \(trainedCaps.count) classifier classes")
    }
    
    func save(classifier: Data, version: Int) throws {
        do {
            try classifier.write(to: classifierFile)
        } catch {
            log("Failed to write classifier: \(error)")
            throw Abort(.internalServerError)
        }
        classifierVersion = version
        log("Updated classifier to version \(version)")
    }

    // MARK: ServerOwner

    let authenticationMethod: PropertyAuthenticationMethod = .accessToken

    func hasReadPermission(for property: UInt32, accessData: Data) -> Bool {
        guard let key = String(data: accessData, encoding: .utf8) else {
            return false
        }
        return writers.contains(key)
    }

    func hasWritePermission(for property: UInt32, accessData: Data) -> Bool {
        guard let key = String(data: accessData, encoding: .utf8) else {
            return false
        }
        return writers.contains(key)
    }

    func hasListAccessPermission(_ accessData: Data) -> Bool {
        guard let key = String(data: accessData, encoding: .utf8) else {
            return false
        }
        return writers.contains(key)
    }

    // MARK: Monitoring

    public let name = "caps"

    private let capCountPropertyId = PropertyId(owner: "caps", uniqueId: 2)

    private let imageCountPropertyId = PropertyId(owner: "caps", uniqueId: 3)

    private let classifierVersionPropertyId = PropertyId(owner: "caps", uniqueId: 4)

    func registerProperties(with monitor: PropertyManager) {
        let capCountProperty = PropertyRegistration(
            uniqueId: capCountPropertyId.uniqueId,
            name: "caps",
            updates: .continuous,
            isLogged: true,
            allowsManualUpdate: false,
            read: { [weak self] in
                return (self?.capCount ?? 0).timestamped()
            })
        monitor.register(capCountProperty, for: self)

        let imageCountProperty = PropertyRegistration(
            uniqueId: imageCountPropertyId.uniqueId,
            name: "images",
            updates: .continuous,
            isLogged: true,
            allowsManualUpdate: false,
            read: { [weak self] in
                return (self?.imageCount ?? 0).timestamped()
            })
        monitor.register(imageCountProperty, for: self)

        let classifierVersionProperty = PropertyRegistration(
            uniqueId: classifierVersionPropertyId.uniqueId,
            name: "classifier",
            updates: .continuous,
            isLogged: true,
            allowsManualUpdate: false,
            read: { [weak self] in
                return (self?.classifierVersion ?? 0).timestamped()
            })
        monitor.register(classifierVersionProperty, for: self)
    }

    private func updateMonitoredPropertiesOnCapChange() {
        try? monitor.logChanged(property: capCountPropertyId, value: capCount.timestamped())
        try? monitor.logChanged(property: imageCountPropertyId, value: imageCount.timestamped())
    }

    private func updateMonitoredClassifierVersionProperty() {
        try? monitor.logChanged(property: classifierVersionPropertyId, value: classifierVersion.timestamped())
    }
}