ChWebsiteApp/CHDataManagement/Generator/ImageGenerator.swift

import Foundation
import AppKit
import SDWebImageAVIFCoder
import SDWebImageWebPCoder

final class ImageGenerator {

    private let storage: Storage

    private let settings: Settings

    private var generatedImages: [String : Set<String>] = [:]

    init(storage: Storage, settings: Settings) {
        self.storage = storage
        self.settings = settings
        self.generatedImages = storage.loadListOfGeneratedImages() ?? [:]
        print("ImageGenerator: Loaded list of \(totalImageCount) already generated images")
    }

    private var outputFolder: String {
        settings.paths.imagesOutputFolderPath
    }

    private var totalImageCount: Int {
        generatedImages.values.reduce(0) { $0 + $1.count }
    }

    @discardableResult
    func save() -> Bool {
        guard storage.save(listOfGeneratedImages: generatedImages) else {
            print("ImageGenerator: Failed to save list of generated images")
            return false
        }
        print("ImageGenerator: Saved list of \(totalImageCount) images")
        return true
    }

    private var avifCommands: Set<String> = []

    func printAvifCommands() {
        avifCommands.sorted().forEach { print($0) }
    }

    /**
     Remove all versions of an image, so that they will be recreated on the next run.

     This function does not remove the images from the output folder.
     */
    func removeVersions(of image: String) {
        generatedImages[image] = nil
        save()
    }

    func recalculateGeneratedImages(by images: Set<String>) {
        self.generatedImages = storage.calculateImages(generatedBy: images, in: outputFolder)
        let versionCount = generatedImages.values.reduce(0) { $0 + $1.count }
        print("Image generator: \(generatedImages.count)/\(images.count) images (\(versionCount) versions)")
    }

    private func hasPreviouslyGenerated(_ version: ImageVersion) -> Bool {
        guard let versions = generatedImages[version.image.id] else {
            return false
        }
        return versions.contains(version.versionId)
    }

    private func needsToGenerate(_ version: ImageVersion) -> Bool {
        if hasPreviouslyGenerated(version) {
            return false
        }
        if exists(version) {
            // Mark as already generated
            hasNowGenerated(version)
            return false
        }
        return true
    }

    private func hasNowGenerated(_ version: ImageVersion) {
        generatedImages[version.image.id, default: []].insert(version.versionId)
    }

    private func removeVersions(for image: String) {
        generatedImages[image] = nil
    }

    // MARK: Files

    private func exists(_ version: ImageVersion) -> Bool {
        storage.hasFileInOutputFolder(version.outputPath)
    }

    private func write(imageData data: Data, of version: ImageVersion) -> Bool {
        return storage.write(data, to: version.outputPath)
    }

    // MARK: Image operations

    func generate(version: ImageVersion) -> Bool {
        guard needsToGenerate(version) else {
            return true
        }

        guard let data = version.image.dataContent() else {
            print("ImageGenerator: Failed to load data for image \(version.image.id)")
            return false
        }

        guard let originalImage = NSImage(data: data) else {
            print("ImageGenerator: Failed to load image \(version.image.id)")
            return false
        }

        let representation = create(image: originalImage, width: CGFloat(version.maximumWidth), height: CGFloat(version.maximumHeight))

        guard let data = create(image: representation, type: version.type, quality: version.quality) else {
            print("ImageGenerator: Failed to get data for type \(version.type) of image \(version.image.id)")
            return false
        }

        if version.type == .avif {
            if version.image.type == .gif {
                // Skip GIFs, since they can't be converted by avifenc
                return true
            }
            // AVIF conversion is very slow, so we save bash commands
            // for the conversion instead
            let baseVersion = ImageVersion(
                image: version.image,
                type: version.image.type,
                maximumWidth: version.maximumWidth,
                maximumHeight: version.maximumHeight)
            let originalImagePath = storage.outputPath(to: baseVersion.outputPath)!.path()
            let generatedImagePath = storage.outputPath(to: version.outputPath)!.path()
            let quality = Int(version.quality * 100)

            avifCommands.insert("avifenc -q \(quality) '\(originalImagePath)' '\(generatedImagePath)'")
            // hasNowGenerated(version)
            return true
        }

        guard write(imageData: data, of: version) else {
            return false
        }
        hasNowGenerated(version)
        return true
    }

    private func create(image originalImage: NSImage, width: CGFloat, height: CGFloat) -> NSBitmapImageRep {
        let sourceRep = originalImage.representations[0]
        let sourceSize = NSSize(width: sourceRep.pixelsWide, height: sourceRep.pixelsHigh)
        let maximumSize = NSSize(width: width, height: height)
        let destinationSize = sourceSize.scaledToFit(in: maximumSize)

        // create NSBitmapRep manually, if using cgImage, the resulting size is wrong
        let representation = 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: representation)
        NSGraphicsContext.saveGraphicsState()
        NSGraphicsContext.current = ctx
        originalImage.draw(in: NSMakeRect(0, 0, destinationSize.width, destinationSize.height))
        ctx?.flushGraphics()
        NSGraphicsContext.restoreGraphicsState()
        return representation
    }

    // MARK: Avif images

    private func create(image: NSBitmapImageRep, type: FileType, quality: CGFloat) -> Data? {
        switch type {
        case .jpg:
            return image.representation(using: .jpeg, properties: [.compressionFactor: NSNumber(value: 0.6)])
        case .png:
            return image.representation(using: .png, properties: [.compressionFactor: NSNumber(value: 0.6)])
        case .avif:
            return createAvif(image: image, quality: 0.7)
        case .webp:
            return createWebp(image: image, quality: 0.8)
        case .gif:
            return image.representation(using: .gif, properties: [.compressionFactor: NSNumber(value: quality)])
        case .svg:
            return nil
        case .tiff:
            return nil
        default:
            return nil
        }
    }

    private func createAvif(image: NSBitmapImageRep, quality: CGFloat) -> Data? {
        return Data()
//        let newImage = NSImage(size: image.size)
//        newImage.addRepresentation(image)
//        return SDImageAVIFCoder.shared.encodedData(with: newImage, format: .AVIF, options: [.encodeCompressionQuality: quality])
    }

    private func createWebp(image: NSBitmapImageRep, quality: CGFloat) -> Data? {
        let newImage = NSImage(size: image.size)
        newImage.addRepresentation(image)
        return SDImageWebPCoder.shared.encodedData(with: newImage, format: .webP, options: [.encodeCompressionQuality: quality])
    }
}