Caps-Server/Sources/App/CapServer.swift

import Foundation
import Vapor
import Clairvoyant

final class CapServer {
    
    private let imageSize = 360
    
    private let thumbnailSize = 100
    
    // MARK: Paths
    
    private let imageFolder: URL

    private let thumbnailFolder: URL
    
    /// The file where the cap count is stored for the grid webpage
    private let gridCountFile: 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 classifierClassesFile: URL
    
    private let changedImagesFile: URL
    
    private let fm = FileManager.default
    
    private let changedImageEntryDateFormatter: DateFormatter

    /// Indicates that the data is loaded
    private(set) var isOperational = false
    
    private(set) var canResizeImages = false
    
    // MARK: Caps

    
    /// The changed images not yet written to disk
    private var unwrittenImageChanges: [(cap: Int, image: Int)] = []
    
    var classifierVersion: Int = 0 {
        didSet {
            writeClassifierVersion()
            Task {
                try? await classifierMetric.update(classifierVersion)
            }
        }
    }

    /**
     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()
            Task {
                _ = try? await capCountMetric.update(caps.count)
                _ = try? await imageCountMetric.update(imageCount)
            }
        }
    }

    var capCount: Int {
        caps.count
    }

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

    init(in folder: URL) {
        self.imageFolder = folder.appendingPathComponent("images")
        self.thumbnailFolder = folder.appendingPathComponent("thumbnails")
        self.gridCountFile = folder.appendingPathComponent("count.js")
        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.classifierClassesFile = folder.appendingPathComponent("classifier.classes")
        self.changedImagesFile = folder.appendingPathComponent("changes.txt")
        self.changedImageEntryDateFormatter = DateFormatter()
        changedImageEntryDateFormatter.dateFormat = "yy-MM-dd-HH-mm-ss"

        // Metric initializers only fail if observer is missing or ID is duplicate
        self.capCountMetric = .init("caps.count",
                                    name: "Number of caps",
                                    description: "The total number of caps in the database")

        self.imageCountMetric = .init("caps.images",
                                      name: "Total images",
                                      description: "The total number of images for all caps")

        self.classifierMetric = .init("caps.classifier",
                                      name: "Classifier Version",
                                      description: "The current version of the image classifier")
    }

    func loadData() throws {
        loadClassifierVersion(at: classifierVersionFile)
        try loadCaps()
        saveCapCountHTML()
        updateGridCapCount()
        try ensureExistenceOfChangedImagesFile()
        organizeImages()
        if let version = getMagickVersion() {
            log("Using ImageMagick \(version.rawValue)")
            canResizeImages = true
        }
        // shrinkImages()
        createMissingThumbnails()
        isOperational = true
    }

    private func loadClassifierVersion(at url: URL) {
        guard exists(url) 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 {
        guard exists(dbFile) else {
            log("No cap database found")
            return
        }
        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
        }
    }

    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;overflow: hidden">
            <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)
    }
    
    private func organizeImages() {
        caps.values.sorted().forEach(organizeImages)
    }

    private func createImageFolder(for cap: Int) throws {
        let folderUrl = folder(of: cap)
        do {
            try fm.createDirectory(at: folderUrl, withIntermediateDirectories: true)
        } catch {
            log("Failed to create folder for cap \(cap): \(error)")
            throw error
        }
    }
    
    /**
     Rearrange images of a cap to ensure that an image exists for each number from 0 to `image count - 1`.
     This is done by using the last images to fill in possible gaps in the sequence.
     E.g. If there are images `0`, `2`, `3`, then `3` will be renamed to `1`.
     - Note: The main image is also changed, if the main image is renamed.
     */
    private func organizeImages(for cap: Cap) {
        var cap = cap
        let folderUrl = folder(of: cap.id)
        guard exists(folderUrl) else {
            try? createImageFolder(for: cap.id)
            cap.count = 0
            caps[cap.id] = cap
            log("Found cap \(cap.id) without image folder")
            return
        }

        guard let images = try? images(in: folderUrl) else {
            log("Failed to get image urls for cap \(cap.id)")
            return
        }
        
        if images.count != cap.count {
            log("\(images.count) instead of \(cap.count) images for cap \(cap.id)")
        }
        
        // Get list of existing images
        var sorted: [(id: Int, url: URL)] = images.compactMap {
            guard let id = Int($0.deletingPathExtension().lastPathComponent.components(separatedBy: "-").last!) else {
                return nil
            }
            return (id, $0)
        }.sorted { $0.id < $1.id }
        
        // Check that all images are available
        for version in 0..<images.count {
            guard version != sorted[version].id else {
                continue
            }
            let lastImage = sorted.popLast()!
            let newUrl = imageUrl(of: cap.id, version: version)
            do {
                try fm.moveItem(at: lastImage.url, to: newUrl)
                log("Moved image \(lastImage.id) to \(version) for cap \(cap.id)")
            } catch {
                log("Failed to move file \(lastImage.url.path) to \(newUrl.path): \(error)")
                return
            }
            if cap.mainImage == lastImage.id {
                cap.mainImage = version
            }
            sorted.insert((version, newUrl), at: version)
        }

        cap.count = sorted.count

        // Fix invalid main image
        if cap.mainImage >= cap.count || cap.mainImage < 0 {
            cap.mainImage = 0
        }

        caps[cap.id] = cap
    }

    // MARK: Paths

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

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

    private func exists(_ url: URL) -> Bool {
        fm.fileExists(atPath: url.path)
    }

    // MARK: Counts

    private func images(in folder: URL) throws -> [URL] {
        do {
            return try fm.contentsOfDirectory(at: folder, includingPropertiesForKeys: nil)
                .filter { $0.pathExtension == "jpg" }
        } catch {
            log("Failed to get image urls for cap \(folder.lastPathComponent): \(error)")
            throw error
        }
    }

    /**
     Get the image count of a cap.
     */
    func count(of cap: Int) throws -> Int {
        let capImageFolder = folder(of: cap)
        guard exists(capImageFolder) else {
            return 0
        }
        return try images(in: capImageFolder).count
    }

    // 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 {
            log("Tried to save image for unknown cap \(cap)")
            throw CapError.unknownId
        }
        var id = 0
        let capFolder = folder(of: cap)
        var capImageUrl = imageUrl(of: cap, version: id)
        if exists(capFolder) {
            while exists(capImageUrl) {
                id += 1
                capImageUrl = imageUrl(of: cap, version: id)
            }
        } else {
            try createImageFolder(for: cap)
        }
        do {
            try data.write(to: capImageUrl)
        } catch {
            log("Failed to write image \(id) for cap \(cap): \(error)")
            throw CapError.invalidFile
        }
        let count = try count(of: cap)
        caps[cap]!.count = count
        addChangedImageToLog(cap: cap, image: id)
        if canResizeImages {
            shrink(imageAt: capImageUrl, size: imageSize, destination: capImageUrl)
            createThumbnail(for: cap)
        }
        log("Added image \(id) for cap \(cap) (\(count) total)")
    }
    
    private func writeChangedImagesToDisk() throws {
        guard !unwrittenImageChanges.isEmpty else {
            return
        }

        try ensureExistenceOfChangedImagesFile()
        
        let handle: FileHandle
        do {
            handle = try FileHandle(forWritingTo: changedImagesFile)
            try handle.seekToEnd()
        } catch {
            log("Failed to open changed images file for writing: \(error)")
            throw error
        }

        var entries = unwrittenImageChanges
        defer {
            unwrittenImageChanges = entries
            try? handle.close()
        }
        let dateString = changedImageEntryDateFormatter.string(from: Date())
        while let entry = entries.popLast() {
            let content = "\(dateString):\(entry.cap):\(entry.image)\n".data(using: .utf8)!
            do {
                try handle.write(contentsOf: content)
            } catch {
                log("Failed to write entry to changed images file: \(error)")
                throw error
            }
        }
    }
    
    private func addChangedImageToLog(cap: Int, image: Int) {
        unwrittenImageChanges.append((cap, image))
        do {
            try writeChangedImagesToDisk()
        } catch {
            log("Failed to save changed image list: \(error)")
        }
    }
    
    private func ensureExistenceOfChangedImagesFile() throws {
        if exists(changedImagesFile) {
            return
        }
        do {
            try Data().write(to: changedImagesFile)
        } catch {
            log("Failed to create changed images file: \(error)")
            throw error
        }
    }
    
    func removeAllEntriesInImageChangeList(before date: Date) {
        guard exists(changedImagesFile) else {
            log("No file for changed images to update")
            return
        }
        do {
            try String(contentsOf: changedImagesFile)
                .components(separatedBy: "\n")
                .filter { $0 != "" }
                .compactMap { line -> String? in
                    guard let entryDate = changedImageEntryDateFormatter.date(from: line.components(separatedBy: ":").first!) else {
                        return nil
                    }
                    guard entryDate > date else {
                        return nil
                    }
                    return line
                }
                .joined(separator: "\n")
                .data(using: .utf8)!
                .write(to: changedImagesFile)
        } catch {
            log("Failed to update changed images file: \(error)")
        }
    }
    
    func switchMainImage(to version: Int, for cap: Int) throws {
        let capImageUrl = imageUrl(of: cap, version: version)
        guard exists(capImageUrl) else {
            log("No image \(version) for cap \(cap)")
            throw CapError.invalidFile
        }
        caps[cap]?.mainImage = version
        if canResizeImages {
            createThumbnail(for: cap)
        }
        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 {
            log("Attempting to add cap \(cap.id) with main image \(cap.mainImage)")
            throw CapError.invalidData
        }
        var cap = cap
        cap.count = 0
        caps[cap.id] = cap
        saveCapCountHTML()
        updateGridCapCount()
        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 capImageUrl = imageUrl(of: existingCap.id, version: cap.mainImage)
        if exists(capImageUrl) {
            updatedCap.mainImage = cap.mainImage
        }
        if let color = cap.color {
            updatedCap.color = color
        }
        caps[existingCap.id] = updatedCap
        log("Updated cap \(existingCap.id)")
    }
    
    func deleteImage(version: Int, for capId: Int) -> Cap? {
        guard let cap = caps[capId] else {
            log("Attempting to delete image \(version) of unknown cap \(capId)")
            return nil
        }
        let capImageUrl = imageUrl(of: capId, version: version)
        guard exists(capImageUrl) else {
            log("Attempting to delete missing image \(version) of cap \(capId)")
            return nil
        }
        organizeImages(for: cap)
        return caps[capId]!
    }

    func delete(cap capId: Int) -> Bool {
        guard caps[capId] != nil else {
            log("Attempting to delete unknown cap \(capId)")
            return false
        }
        // 1. Remove all images
        do {
            let imageFolderUrl = folder(of: capId)
            if exists(imageFolderUrl) {
                try fm.removeItem(at: imageFolderUrl)
            }
        } catch {
            log("Failed to delete image folder of cap \(capId): \(error)")
            return false
        }
        // 2. Remove thumbnail
        do {
            let url = thumbnail(of: capId)
            if exists(url) {
                try fm.removeItem(at: url)
            }
        } catch {
            log("Failed to delete thumbnail of cap \(capId): \(error)")
            return false
        }

        // 3. Remove cap
        caps[capId] = nil
        saveCapCountHTML()
        updateGridCapCount()
        return true
    }
    
    // MARK: Classifier
    
    func saveTrainedClasses(content: String) throws {
        let classes = content.components(separatedBy: ",")
        
        // Validate input
        try classes.forEach { s in
            guard let id = Int(s) else {
                log("Invalid id '\(s)' in uploaded id list")
                throw Abort(.badRequest)
            }
            guard caps[id] != nil else {
                log("Unknown id '\(id)' in uploaded id list")
                throw Abort(.badRequest)
            }
        }
        
        guard let data = content.data(using: .utf8) else {
            log("Failed to get classes data for writing")
            throw Abort(.internalServerError)
        }
        
        do {
            try data.write(to: classifierClassesFile)
            log("Updated \(classes.count) classifier classes")
        } catch {
            log("Failed to write classifier classes: \(error)")
            throw Abort(.internalServerError)
        }
    }
    
    func save(classifier: Data, version: Int) throws {
        do {
            try classifier.write(to: classifierFile)
        } catch {
            log("Failed to write classifier \(version): \(error)")
            throw Abort(.internalServerError)
        }
        classifierVersion = version
        log("Updated classifier to version \(version)")
    }
    
    // MARK: Grid
    
    func getListOfMissingThumbnails() -> [Int] {
        caps.keys.filter { !exists(thumbnail(of: $0)) }
    }
    
    func saveThumbnail(_ data: Data, for cap: Int) {
        let url = thumbnail(of: cap)
        do {
            try data.write(to: url)
        } catch {
            log("Failed to save thumbnail \(cap): \(error)")
        }
    }
    
    private func updateGridCapCount() {
        do {
            try "const numberOfCaps = \(capCount);"
                .data(using: .utf8)!
                .write(to: gridCountFile)
        } catch {
            log("Failed to save grid cap count: \(error)")
        }
    }
    
    func createMissingThumbnails() {
        let thumbnailsToCreate = getListOfMissingThumbnails()
        guard !thumbnailsToCreate.isEmpty else {
            return
        }
        guard canResizeImages else {
            log("Can't create thumbnails, missing ImageMagick")
            return
        }
        log("Creating \(thumbnailsToCreate.count) thumbnails")
        for cap in thumbnailsToCreate {
            createThumbnail(for: cap)
        }
    }
    
    func createThumbnail(for cap: Int) {
        guard let version = caps[cap]?.mainImage else {
            return
        }
        let mainImageUrl = imageUrl(of: cap, version: version)
        let thumbnailUrl = thumbnail(of: cap)
        shrink(imageAt: mainImageUrl, size: thumbnailSize, destination: thumbnailUrl)
    }

    // MARK: Monitoring

    private let capCountMetric: Metric<Int>

    private let imageCountMetric: Metric<Int>

    private let classifierMetric: Metric<Int>
    
    // MARK: Maintenance
    
    private func getMagickVersion() -> SemanticVersion? {
        do {
            let command = "convert -version"
            let (code, output) = try safeShell(command)
            guard code == 0,
                  let line = output.components(separatedBy: "\n").first,
                  line.hasPrefix("Version: ImageMagick ") else {
                log("Missing dependency ImageMagick: " + output)
                return nil
            }
            guard let versionString = line
                .replacingOccurrences(of: "Version: ImageMagick ", with: "")
                .components(separatedBy: "-").first  else {
                log("Invalid ImageMagick version: " + output)
                return nil
            }
            guard let version = SemanticVersion(rawValue: versionString) else {
                log("Invalid ImageMagick version: " + output)
                return nil
            }
            return version
        } catch {
            log("Failed to check dependency ImageMagick: \(error)")
            return nil
        }
    }
    
    func shrinkImages() {
        guard canResizeImages else {
            log("Can't resize images, missing ImageMagick")
            return
        }
        let imageFolders: [URL]
        do {
            imageFolders = try fm.contentsOfDirectory(at: imageFolder, includingPropertiesForKeys: nil)
        } catch {
            log("Failed to get all image folders")
            return
        }
        for folder in imageFolders {
            guard let images = try? self.images(in: folder) else {
                continue
            }
            for imageUrl in images {
                shrink(imageAt: imageUrl, size: imageSize, destination: imageUrl)
            }
        }
    }
    
    private func shrink(imageAt url: URL, size: Int, destination: URL) {
        do {
            let command = "convert \(url.path) -resize '\(size)x\(size)>' \(destination.path)"
            let (code, output) = try safeShell(command)
            if code != 0 {
                log("Failed to shrink image \(url.path): " + output)
            }
        } catch {
            log("Failed to shrink image \(url.path): \(error)")
        }
    }
    
    private func safeShell(_ command: String) throws -> (code: Int32, output: String) {
        let task = Process()
        let pipe = Pipe()

        task.standardOutput = pipe
        task.standardError = pipe
        task.arguments = ["-cl", command]
        task.executableURL = URL(fileURLWithPath: "/bin/bash")
        task.standardInput = nil

        try task.run()
        task.waitUntilExit()
        
        let data = pipe.fileHandleForReading.readDataToEndOfFile()
        let output = String(data: data, encoding: .utf8)!
        return (task.terminationStatus, output)
    }
}