CHGenerator/Sources/Generator/Processing/ImageGenerator.swift

import Foundation
import AppKit
import CryptoKit
import Darwin.C

final class ImageGenerator {

    private let imageOptimSupportedFileExtensions: Set<String> = ["jpg", "png", "svg"]

    private let imageOptimizationBatchSize = 50

    /**
     The path to the input folder.
     */
    private let input: URL

    /**
     The path to the output folder
     */
    private let output: URL

    private let imageReader: ImageReader
    
    /**
     All warnings produced for images during generation
     */
    private var imageWarnings: [String]

    /// The images which could not be found, but are required for the site (`key`: image path, `value`: source element path)
    private var missingImages: [String : String]

    /// Images which could not be read (`key`: file path relative to content, `value`: source element path)
    private var unreadableImages: [String : String]

    private var unhandledImages: [String: String] = [:]

    /// All images modified or created during this generator run.
    private var generatedImages: Set<String> = []

    /// The images optimized by ImageOptim
    private var optimizedImages: Set<String> = []

    private var failedImages: [(path: String, message: String)] = []

    private var numberOfGeneratedImages = 0

    private let numberOfTotalImages: Int

    private lazy var numberOfImagesToCreate = jobs.reduce(0) { $0 + $1.images.count } + multiJobs.count * 2

    private var numberOfImagesToOptimize = 0

    private var numberOfOptimizedImages = 0

    private let images: ImageData

    private lazy var jobs: [(source: String, images: [ImageJob])] = images.jobs
        .sorted { $0.key < $1.key }
        .map { (source: $0.key, images: $0.value) }
        .filter {
            // Only load image if required
            let imageHasChanged = imageReader.imageHasChanged(at: $0.source)
            return imageHasChanged || $0.images.contains { job in
                job.alwaysGenerate || !output.appendingPathComponent(job.destination).exists
            }
        }

    private lazy var multiJobs: [String : String] = {
        let imagesToGenerate: Set<String> = jobs.reduce([]) { $0.union($1.images.map { $0.destination }) }
        return images.multiJobs.filter { imagesToGenerate.contains($0.key) }
    }()

    init(input: URL, output: URL, reader: ImageReader, images: ImageData) {
        self.input = input
        self.output = output
        self.imageReader = reader
        self.images = images
        self.numberOfTotalImages = images.jobs.reduce(0) { $0 + $1.value.count }
            + images.multiJobs.count * 2
        self.imageWarnings = images.warnings
        self.missingImages = images.missing
        self.unreadableImages = images.unreadable
    }

    func generateImages() {
        var notes: [String] = []
        func addIfNotZero(_ count: Int, _ name: String) {
            guard count > 0 else {
                return
            }
            notes.append("\(count) \(name)")
        }

        print("       Changed sources: \(jobs.count)/\(images.jobs.count)")
        print("          Total images: \(numberOfTotalImages) (\(numberOfTotalImages - imageReader.numberOfFilesAccessed) versions)")

        for (source, jobs) in jobs {
            create(images: jobs, from: source)
        }
        for (baseImage, source) in multiJobs {
            createMultiImages(from: baseImage, path: source)
        }
        print("      Generated images: \(numberOfGeneratedImages)/\(numberOfImagesToCreate)")
        optimizeImages()
        print("      Optimized images: \(numberOfOptimizedImages)/\(numberOfImagesToOptimize)")

        addIfNotZero(missingImages.count, "missing images")
        addIfNotZero(unreadableImages.count, "unreadable images")
        print("              Warnings: \(imageWarnings.count)")
        if !notes.isEmpty {
            print("                 Notes: " + notes.joined(separator: ", "))
        }
    }

    private func create(images: [ImageJob], from source: String) {
        guard let image = imageReader.getImage(atPath: source) else {
            unreadableImages[source] = images.first!.destination
            didGenerateImage(count: images.count)
            return
        }
        let jobs = imageReader.imageHasChanged(at: source) ? images : images.filter(isMissing)
        // Update all images
        jobs.forEach { job in
            // Prevent memory overflow due to repeated NSImage operations
            autoreleasepool {
                create(job: job, from: image, source: source)
                didGenerateImage()
            }
        }
    }

    private func isMissing(_ job: ImageJob) -> Bool {
        job.alwaysGenerate || !output.appendingPathComponent(job.destination).exists
    }


    private func create(job: ImageJob, from image: NSImage, source: String) {
        let destinationUrl = output.appendingPathComponent(job.destination)
        create(job: job, from: image, source: source, at: destinationUrl)
    }

    private func create(job: ImageJob, from image: NSImage, source: String, at destinationUrl: URL) {
        // Ensure that image file is supported
        let ext = destinationUrl.pathExtension.lowercased()
        guard ImageType(fileExtension: ext) != nil else {
            fatalError()
        }

        let destinationExtension = destinationUrl.pathExtension.lowercased()
        guard let type = ImageType(fileExtension: destinationExtension)?.fileType else {
            unhandledImages[source] = job.destination
            return
        }

        let desiredWidth = CGFloat(job.width)

        let sourceRep = image.representations[0]
        let destinationSize = NSSize(width: sourceRep.pixelsWide, height: sourceRep.pixelsHigh)
            .scaledDown(to: desiredWidth)

        // create NSBitmapRep manually, if using cgImage, the resulting size is wrong
        let rep = NSBitmapImageRep(bitmapDataPlanes: nil,
                                   pixelsWide: Int(destinationSize.width),
                                   pixelsHigh: Int(destinationSize.height),
                                   bitsPerSample: 8,
                                   samplesPerPixel: 4,
                                   hasAlpha: true,
                                   isPlanar: false,
                                   colorSpaceName: NSColorSpaceName.deviceRGB,
                                   bytesPerRow: Int(destinationSize.width) * 4,
                                   bitsPerPixel: 32)!

        let ctx = NSGraphicsContext(bitmapImageRep: rep)
        NSGraphicsContext.saveGraphicsState()
        NSGraphicsContext.current = ctx
        image.draw(in: NSMakeRect(0, 0, destinationSize.width, destinationSize.height))
        ctx?.flushGraphics()
        NSGraphicsContext.restoreGraphicsState()

        // Get NSData, and save it
        guard let data = rep.representation(using: type, properties: [.compressionFactor: NSNumber(value: job.quality)]) else {
            markImageAsFailed(source, error: "Failed to get data")
            return
        }
        do {
            try data.createFolderAndWrite(to: destinationUrl)
        } catch {
            markImageAsFailed(job.destination, error: "Failed to write image (\(error))")
            return
        }
        generatedImages.insert(job.destination)
    }

    private func markImageAsFailed(_ source: String, error: String) {
        failedImages.append((source, error))
    }

    private func createMultiImages(from source: String, path: String) {
        guard generatedImages.contains(source) else {
            didGenerateImage(count: 2)
            return
        }

        let sourceUrl = output.appendingPathComponent(source)
        let sourcePath = sourceUrl.path
        guard sourceUrl.exists else {
            missingImages[source] = path
            didGenerateImage(count: 2)
            return
        }

        let avifPath = source.dropAfterLast(".") + ".avif"
        createAVIF(at: output.appendingPathComponent(avifPath).path, from: sourcePath)
        generatedImages.insert(avifPath)
        didGenerateImage()

        let webpPath = source.dropAfterLast(".") + ".webp"
        createWEBP(at: output.appendingPathComponent(webpPath).path, from: sourcePath)
        generatedImages.insert(webpPath)
        didGenerateImage()

        compress(at: sourcePath)
    }

    private func createAVIF(at destination: String, from source: String, quality: Int = 55, effort: Int = 5) {
        let folder = destination.dropAfterLast("/")
        let command = "npx avif --input=\(source) --quality=\(quality) --effort=\(effort) --output=\(folder) --overwrite"
        do {
            let output = try safeShell(command)
            if output == "" {
                return
            }
            markImageAsFailed(destination, error: "Failed to create AVIF image: \(output)")
        } catch {
            markImageAsFailed(destination, error: "Failed to create AVIF image")
        }
    }

    private func createWEBP(at destination: String, from source: String, quality: Int = 75) {
        let command = "cwebp \(source) -q \(quality) -o \(destination)"
        do {
            let output = try safeShell(command)
            if !output.contains("Error") {
                return
            }
            markImageAsFailed(destination, error: "Failed to create WEBP image: \(output)")
        } catch {
            markImageAsFailed(destination, error: "Failed to create WEBP image: \(error)")
        }
    }

    private func compress(at destination: String, quality: Int = 70) {
        let command = "magick convert \(destination) -quality \(quality)% \(destination)"
        do {
            let output = try safeShell(command)
            if output == "" {
                return
            }
            markImageAsFailed(destination, error: "Failed to compress image: \(output)")
        } catch {
            markImageAsFailed(destination, error: "Failed to compress image: \(error)")
        }
    }

    private func optimizeImages() {
        let all = generatedImages
            .filter { imageOptimSupportedFileExtensions.contains($0.lastComponentAfter(".")) }
            .map { output.appendingPathComponent($0).path }
        for i in stride(from: 0, to: all.count, by: imageOptimizationBatchSize) {
            let endIndex = min(i+imageOptimizationBatchSize, all.count)
            let batch = all[i..<endIndex]
            if optimizeImageBatch(batch) {
                optimizedImages.formUnion(batch)
            }
            didOptimizeImage(count: batch.count)
        }
    }

    private func optimizeImageBatch(_ batch: ArraySlice<String>) -> Bool {
        let command = "imageoptim " + batch.joined(separator: " ")
        do {
            let output = try safeShell(command)
            if output.contains("Finished") {
                return true
            }

            for image in batch {
                markImageAsFailed(image, error: "Failed to optimize image: \(output)")
            }
            return true
        } catch {
            for image in batch {
                markImageAsFailed(image, error: "Failed to optimize image: \(error)")
            }
            return false
        }
    }

    // MARK: Output

    private func didGenerateImage(count: Int = 1) {
        numberOfGeneratedImages += count
        print("      Generated images: \(numberOfGeneratedImages)/\(numberOfImagesToCreate)    \r", terminator: "")
        fflush(stdout)
    }

    private func didOptimizeImage(count: Int) {
        numberOfOptimizedImages += count
        print("      Optimized images: \(numberOfOptimizedImages)/\(numberOfImagesToOptimize)    \r", terminator: "")
        fflush(stdout)
    }

    func writeResults(to file: URL) {
        var lines: [String] = []
        func add<S>(_ name: String, _ property: S, convert: (S.Element) -> String) where S: Sequence {
            let elements = property.map { "  " + convert($0) }.sorted()
            guard !elements.isEmpty else {
                return
            }
            lines.append("\(name):")
            lines.append(contentsOf: elements)
        }
        add("Missing images", missingImages) { "\($0.key) (required by \($0.value))" }
        add("Unreadable images", unreadableImages) { "\($0.key) (required by \($0.value))" }
        add("Failed images", failedImages) { "\($0.path): \($0.message)" }
        add("Unhandled images", unhandledImages) { "\($0.value) (from \($0.key))" }
        add("Warnings", imageWarnings) { $0 }
        add("Generated images", generatedImages) { $0 }
        add("Optimized images", optimizedImages) { $0 }
        let data = lines.joined(separator: "\n").data(using: .utf8)!
        do {
            try data.createFolderAndWrite(to: file)
        } catch {
            print("    Failed to save log: \(error)")
        }
    }
}