CHGenerator/Sources/Generator/Files/ImageGenerator.swift

import Foundation
import AppKit
import CryptoKit

private struct ImageJob {

    let destination: String

    let width: Int

    let path: String
}

final class ImageGenerator {

    /**
     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 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> = []

    /**
     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, width: Int, height: Int?) -> 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)

        guard let existingSource = imageJobs[source] else {
            imageJobs[source] = [job]
            return scaledSize
        }

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

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

    func createImages() {
        for (source, jobs) in imageJobs.sorted(by: { $0.key < $1.key }) {
            create(images: jobs, from: source)
        }
        printMissingImages()
        printImageWarnings()
        printGeneratedImages()
        printTotalImageCount()
    }

    private func printMissingImages() {
        guard !missingImages.isEmpty else {
            return
        }
        print("\(missingImages.count) missing images:")
        for (source, path) in missingImages {
            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 {
        !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)

        // 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(0.7)]) 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)
    }
}