CHGenerator/Sources/Generator/Files/ImageGenerator.swift

import Foundation
import AppKit
import CryptoKit
import Darwin.C

private struct ImageJob {

    let destination: String

    let width: Int

    let path: String

    let quality: Float

    let alwaysGenerate: Bool
}

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

    /**
     The images to generate.

     The key is the image source path relative to the input folder, and the values are the destination path (relative to the output folder) and the required image width.
     */
    private var imageJobs: [String : [ImageJob]] = [:]

    /**
     The images for which to generate multiple versions

     The key is the source file, the value is the path of the requiring page.
     */
    private var multiImageJobs: [String : String] = [:]

    /**
     The images which could not be found, but are required for the site.

     The key is the image path, and the value is the page that requires it.
     */
    private var missingImages: [String : String] = [:]

    /**
     All warnings produced for images during generation
     */
    private var imageWarnings: Set<String> = []

    /**
     All images required by the site.

     The values are the destination paths of the images, relative to the output folder
     */
    private var requiredImages: Set<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> = []

    /**
     A cache to get the size of source images, so that files don't have to be loaded multiple times.

     The key is the absolute source path, and the value is the image size
     */
    private var imageSizeCache: [String : NSSize] = [:]

    private var fileUpdates: FileUpdateChecker

    init(input: URL, output: URL) {
        self.fileUpdates = FileUpdateChecker(input: input)
        self.input = input
        self.output = output
    }

    func writeDetectedFileChangesToDisk() {
        fileUpdates.writeDetectedFileChangesToDisk()
    }

    private func getImageSize(atPath path: String) -> NSSize? {
        if let size = imageSizeCache[path] {
            return size
        }
        guard let image = getImage(atPath: path) else {
            return nil
        }
        let size = image.size
        imageSizeCache[path] = size
        return size
    }

    private func getImage(atPath path: String) -> NSImage? {
        guard let data = getData(atPath: path) else {
            log.add(error: "Failed to load file", source: path)
            return nil
        }
        guard let image = NSImage(data: data) else {
            log.add(error: "Failed to read image", source: path)
            return nil
        }
        return image
    }

    private func getData(atPath path: String) -> Data? {
        let url = input.appendingPathComponent(path)
        guard url.exists else {
            return nil
        }
        do {
            let data = try Data(contentsOf: url)
            fileUpdates.didLoad(data, at: path)
            return data
        } catch {
            log.add(error: "Failed to read data", source: path, error: error)
            return nil
        }
    }

    func requireImage(at destination: String, generatedFrom source: String, requiredBy path: String, quality: Float, width: Int, height: Int?, alwaysGenerate: Bool) -> NSSize {
        requiredImages.insert(destination)
        let height = height.unwrapped(CGFloat.init)
        let sourceUrl = input.appendingPathComponent(source)
        guard sourceUrl.exists else {
            missingImages[source] = path
            return .zero
        }

        guard let imageSize = getImageSize(atPath: source) else {
            missingImages[source] = path
            return .zero
        }
        let scaledSize = imageSize.scaledDown(to: CGFloat(width))

        // Check desired height, then we can forget about it
        if let height = height {
            let expectedHeight = scaledSize.width / CGFloat(width) * height
            if abs(expectedHeight - scaledSize.height) > 2 {
                addWarning("Invalid height (\(scaledSize.height) instead of \(expectedHeight))", destination: destination, path: path)
            }
        }

        let job = ImageJob(
            destination: destination,
            width: width,
            path: path,
            quality: quality,
            alwaysGenerate: alwaysGenerate)
        insert(job: job, source: source)

        return scaledSize
    }

    private func insert(job: ImageJob, source: String) {
        guard let existingSource = imageJobs[source] else {
            imageJobs[source] = [job]
            return
        }

        guard let existingJob = existingSource.first(where: { $0.destination == job.destination }) else {
            imageJobs[source] = existingSource + [job]
            return
        }

        if existingJob.width != job.width {
            addWarning("Multiple image widths (\(existingJob.width) and \(job.width))", destination: job.destination, path: "\(existingJob.path) and \(job.path)")
        }
    }

    func createImages() {
        var count = 0
        for (source, jobs) in imageJobs.sorted(by: { $0.key < $1.key }) {
            print(String(format: "Creating images: %4d / %d\r", count, imageJobs.count), terminator: "")
            fflush(stdout)
            create(images: jobs, from: source)
            count += 1
        }
        print("                                     \r", terminator: "")
        createMultiImages()
        optimizeImages()
        printMissingImages()
        printImageWarnings()
        printGeneratedImages()
        printTotalImageCount()
    }

    private func printMissingImages() {
        guard !missingImages.isEmpty else {
            return
        }
        print("\(missingImages.count) missing images:")
        let sort = missingImages.sorted { (a, b) in
            a.value < b.value && a.key < b.key
        }
        for (source, path) in sort {
            print("  \(source) (required by \(path))")
        }
    }

    private func printImageWarnings() {
        guard !imageWarnings.isEmpty else {
            return
        }
        print("\(imageWarnings.count) image warnings:")
        for imageWarning in imageWarnings {
            print(imageWarning)
        }
    }

    private func printGeneratedImages() {
        guard !generatedImages.isEmpty else {
            return
        }
        print("\(generatedImages.count) images generated:")
        for image in generatedImages {
            print("  " + image)
        }
    }

    private func printTotalImageCount() {
        print("\(requiredImages.count) images")
    }

    private func addWarning(_ message: String, destination: String, path: String) {
        let warning = "  \(destination): \(message) required by \(path)"
        imageWarnings.insert(warning)
    }

    private func addWarning(_ message: String, job: ImageJob) {
        addWarning(message, destination: job.destination, path: job.path)
    }

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

    private func create(images: [ImageJob], from source: String) {
        // Only load image if required
        let imageHasChanged = fileUpdates.fileHasChanged(at: source)
        guard imageHasChanged || images.contains(where: isMissing) else {
            return
        }

        guard let image = getImage(atPath: source) else {
            missingImages[source] = images.first?.path
            return
        }
        let jobs = imageHasChanged ? 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)
            }
        }
    }

    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 {
            addWarning("Invalid image extension \(destinationExtension)", job: job)
            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 {
            addWarning("Failed to get data", job: job)
            return
        }
        do {
            try data.createFolderAndWrite(to: destinationUrl)
        } catch {
            addWarning("Failed to write image (\(error))", job: job)
            return
        }
        generatedImages.insert(job.destination)
    }

    /**
     Create images of different types.

     This function generates versions for the given image, including png/jpg, avif, and webp. Different pixel density versions (1x and 2x) are also generated.
     - Parameter destination: The path to the destination file
     */
    @discardableResult
    func requireMultiVersionImage(source: String, destination: String, requiredBy path: String, width: Int, desiredHeight: Int?) -> NSSize {
        // Add @1x version
        _ = requireScaledMultiImage(source: source, destination: destination, requiredBy: path, width: width, desiredHeight: desiredHeight)

        // Add @2x version
        return requireScaledMultiImage(
            source: source,
            destination: destination.insert("@2x", beforeLast: "."),
            requiredBy: path,
            width: width * 2,
            desiredHeight: desiredHeight.unwrapped { $0 * 2 })
    }

    @discardableResult
    private func requireScaledMultiImage(source: String, destination: String, requiredBy path: String, width: Int, desiredHeight: Int?) -> NSSize {
        let rawDestinationPath = destination.dropAfterLast(".")
        let avifPath = rawDestinationPath + ".avif"
        let webpPath = rawDestinationPath + ".webp"
        let needsGeneration = !output.appendingPathComponent(avifPath).exists || !output.appendingPathComponent(webpPath).exists

        let size = requireImage(at: destination, generatedFrom: source, requiredBy: path, quality: 1.0, width: width, height: desiredHeight, alwaysGenerate: needsGeneration)
        multiImageJobs[destination] = path
        return size
    }

    private func createMultiImages() {
        let sort = multiImageJobs.sorted { $0.value < $1.value && $0.key < $1.key }
        var count = 1
        for (baseImage, path) in sort {
            print(String(format: "Creating image versions: %4d / %d\r", count, sort.count), terminator: "")
            fflush(stdout)
            createMultiImages(from: baseImage, path: path)
            count += 1
        }
        print("                                     \r", terminator: "")
    }

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

        let sourceUrl = output.appendingPathComponent(source)
        let sourcePath = sourceUrl.path
        guard sourceUrl.exists else {
            addWarning("No image at path \(sourcePath)", destination: source, path: path)
            missingImages[source] = path
            return
        }

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

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

        compress(at: source)
    }

    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 {
            _ = try FileSystem.safeShell(command)
        } catch {
            addWarning("Failed to create AVIF image", destination: destination, path: destination)
        }
    }

    private func createWEBP(at destination: String, from source: String, quality: Int = 75) {
        let command = "cwebp \(source) -q \(quality) -o \(destination)"
        do {
            _ = try FileSystem.safeShell(command)
        } catch {
            addWarning("Failed to create WEBP image", destination: destination, path: destination)
        }
    }

    private func compress(at destination: String, quality: Int = 70) {
        let command = "magick convert \(destination) -quality \(quality)% \(destination)"
        do {
            _ = try FileSystem.safeShell(command)
        } catch {
            addWarning("Failed to compress image", destination: destination, path: destination)
        }
    }

    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]
            print(String(format: "Optimizing images: %4d / %d\r", endIndex, all.count), terminator: "")
            fflush(stdout)
            if optimizeImageBatch(batch) {
                optimizedImages.formUnion(batch)
            }
        }
        print("                                     \r", terminator: "")
        fflush(stdout)
        print("\(optimizedImages.count) images optimized")
    }

    private func optimizeImageBatch(_ batch: ArraySlice<String>) -> Bool {
        let command = "imageoptim " + batch.joined(separator: " ")
        do {
            _ = try FileSystem.safeShell(command)
            return true
        } catch {
            addWarning("Failed to optimize images", destination: "", path: "")
            return false
        }
    }
}