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: 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> = []

    private var numberOfGeneratedImages = 0

    private let numberOfTotalImages: Int

    private lazy var numberImagesToCreate = 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
    }

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

        print("       Changed sources: \(jobs.count)/\(images.jobs.count)")
        print("          Total images: \(numberOfTotalImages) (\(numberOfTotalImages - imageReader.numberOfFilesAccessed) versions)")
        print("              Warnings: \(images.warnings.count)")
        if !notes.isEmpty {
            print("                 Notes: " + notes.joined(separator: ", "))
        }
        for (source, jobs) in jobs {
            create(images: jobs, from: source)
        }
        for (baseImage, source) in multiJobs {
            createMultiImages(from: baseImage, path: source)
        }
        print("      Generated images: \(numberOfGeneratedImages)/\(numberImagesToCreate)")
        optimizeImages()
        print("      Optimized images: \(numberOfOptimizedImages)/\(numberOfImagesToOptimize)")
    }

    private func create(images: [ImageJob], from source: String) {
        guard let image = imageReader.getImage(atPath: source) else {
            // TODO: Add to failed images
            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 {
            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)
    }

    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 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 {
            addWarning("No image at path \(sourcePath)", destination: source, path: 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: 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 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 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 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]
            if optimizeImageBatch(batch) {
                optimizedImages.formUnion(batch)
            }
            didOptimizeImage(count: batch.count)
        }
    }

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

    // MARK: Output

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

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