konverter/Sources/Konverter/Convert/PDFConverter.swift

import AppKit
import CoreGraphics
import CoreText
import Foundation
import PDFKit

/// Everything PDF: rasterising pages, building PDFs from images, text
/// extraction, splitting, merging, rotating and password handling.
enum PDFConverter {
    static let defaultDPI: CGFloat = 300

    // MARK: - Images → PDF

    static func imageToPDF(_ source: URL, at destination: URL, cancellation: Cancellation? = nil) throws {
        guard let image = Raster.firstImage(at: source) else {
            throw ConversionError.unreadable(source.lastPathComponent)
        }
        try imagesToPDF([(image, nil)], at: destination, cancellation: cancellation)
    }

    /// Each image becomes one page, sized so the image keeps its aspect ratio
    /// at 72 dpi (a 1440×900 screenshot becomes a 1440×900 pt page).
    static func imagesToPDF(
        _ pages: [(image: CGImage, name: String?)],
        at destination: URL,
        pageSize: CGSize? = nil,
        cancellation: Cancellation? = nil
    ) throws {
        guard !pages.isEmpty else { throw ConversionError.failed("Nothing to put in the PDF.") }

        var mediaBox = CGRect(origin: .zero, size: pages[0].name == nil
            ? CGSize(width: pages[0].image.width, height: pages[0].image.height)
            : CGSize(width: 595, height: 842))
        if let pageSize { mediaBox = CGRect(origin: .zero, size: pageSize) }

        guard let context = CGContext(destination as CFURL, mediaBox: &mediaBox, nil) else {
            throw ConversionError.failed("Could not create \(destination.lastPathComponent).")
        }

        for page in pages {
            try cancellation?.check()
            let image = page.image
            var box = mediaBox
            if page.name == nil || pageSize == nil {
                box = CGRect(x: 0, y: 0, width: image.width, height: image.height)
            }
            context.beginPDFPage([kCGPDFContextMediaBox: CGRectToData(box)] as CFDictionary)

            let scale = min(box.width / CGFloat(image.width), box.height / CGFloat(image.height))
            let size = CGSize(width: CGFloat(image.width) * scale, height: CGFloat(image.height) * scale)
            let rect = CGRect(
                x: (box.width - size.width) / 2,
                y: (box.height - size.height) / 2,
                width: size.width,
                height: size.height
            )
            context.draw(image, in: rect)
            context.endPDFPage()
        }
        context.closePDF()
    }

    private static func CGRectToData(_ rect: CGRect) -> CFData {
        var rect = rect
        return CFDataCreate(nil, withUnsafeBytes(of: &rect) { $0.bindMemory(to: UInt8.self).baseAddress }, MemoryLayout<CGRect>.size)
    }

    // MARK: - PDF → images

    /// Renders each page to `format` at `dpi`. Returns the files written.
    @discardableResult
    static func renderPages(
        _ source: URL,
        to format: Format,
        into directory: URL,
        dpi: CGFloat = 300,
        pageNumbers: [Int]? = nil,
        quality: Double = 0.85,
        cancellation: Cancellation? = nil,
        progress: ((Double) -> Void)? = nil
    ) throws -> [URL] {
        guard let document = PDFDocument(url: source) else {
            throw ConversionError.unreadable(source.lastPathComponent)
        }
        let base = source.deletingPathExtension().lastPathComponent
        try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)

        let indices = pageNumbers ?? Array(0..<document.pageCount)
        var written: [URL] = []
        let scale = dpi / 72
        let padding = String(indices.count).count

        for (offset, pageIndex) in indices.enumerated() {
            try cancellation?.check()
            guard let page = document.page(at: pageIndex) else { continue }
            let bounds = page.bounds(for: .mediaBox)
            let pixelWidth = Int((bounds.width * scale).rounded())
            let pixelHeight = Int((bounds.height * scale).rounded())
            guard pixelWidth > 0, pixelHeight > 0 else { continue }

            guard let context = CGContext(
                data: nil,
                width: pixelWidth,
                height: pixelHeight,
                bitsPerComponent: 8,
                bytesPerRow: 0,
                space: CGColorSpace(name: CGColorSpace.sRGB) ?? CGColorSpaceCreateDeviceRGB(),
                bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue
            ) else { continue }

            context.setFillColor(CGColor(gray: 1, alpha: 1))
            context.fill(CGRect(x: 0, y: 0, width: pixelWidth, height: pixelHeight))
            context.scaleBy(x: scale, y: scale)
            context.translateBy(x: -bounds.origin.x, y: -bounds.origin.y)
            page.draw(with: .mediaBox, to: context)

            guard let image = context.makeImage() else { continue }
            let number = String(pageIndex + 1).leftPadded(to: padding)
            let target = OutputNaming.availableURL(
                for: directory.appendingPathComponent("\(base)-\(number)").appendingPathExtension(format.ext)
            )
            try Raster.write(image, to: target, format: format, quality: quality, metadata: nil)
            written.append(target)
            progress?(Double(offset + 1) / Double(indices.count))
        }
        return written
    }

    // MARK: - Text

    static func extractText(_ source: URL, pageBreak: String = "\n\n") -> String {
        guard let document = PDFDocument(url: source) else { return "" }
        return (0..<document.pageCount)
            .compactMap { document.page(at: $0)?.string }
            .joined(separator: pageBreak)
    }

    // MARK: - Document surgery

    static func merge(_ sources: [URL], at destination: URL, cancellation: Cancellation? = nil) throws {
        let merged = PDFDocument()
        var index = 0
        for source in sources {
            try cancellation?.check()
            guard let document = PDFDocument(url: source) else {
                throw ConversionError.unreadable(source.lastPathComponent)
            }
            for pageIndex in 0..<document.pageCount {
                try cancellation?.check()
                guard let page = document.page(at: pageIndex)?.copy() as? PDFPage else { continue }
                merged.insert(page, at: index)
                index += 1
            }
        }
        guard merged.pageCount > 0 else { throw ConversionError.failed("No pages to merge.") }
        guard merged.write(to: destination) else {
            throw ConversionError.failed("Could not write \(destination.lastPathComponent).")
        }
    }

    /// Splits into chunks of `pagesPerFile`, writing into a sibling folder.
    static func split(
        _ source: URL,
        pagesPerFile: Int,
        into directory: URL,
        cancellation: Cancellation? = nil
    ) throws -> [URL] {
        guard let document = PDFDocument(url: source) else {
            throw ConversionError.unreadable(source.lastPathComponent)
        }
        let base = source.deletingPathExtension().lastPathComponent
        try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)

        var written: [URL] = []
        var start = 0
        var part = 1
        while start < document.pageCount {
            try cancellation?.check()
            let chunk = PDFDocument()
            var cursor = 0
            for pageIndex in start..<min(start + pagesPerFile, document.pageCount) {
                guard let page = document.page(at: pageIndex)?.copy() as? PDFPage else { continue }
                chunk.insert(page, at: cursor)
                cursor += 1
            }
            let target = OutputNaming.availableURL(
                for: directory.appendingPathComponent("\(base)-\(String(format: "%02d", part)).pdf")
            )
            guard chunk.write(to: target) else {
                throw ConversionError.failed("Could not write \(target.lastPathComponent).")
            }
            written.append(target)
            start += pagesPerFile
            part += 1
        }
        return written
    }

    static func rotate(_ source: URL, degrees: Int, at destination: URL) throws {
        guard let document = PDFDocument(url: source) else {
            throw ConversionError.unreadable(source.lastPathComponent)
        }
        for index in 0..<document.pageCount {
            guard let page = document.page(at: index) else { continue }
            page.rotation = (page.rotation + degrees) % 360
        }
        guard document.write(to: destination) else {
            throw ConversionError.failed("Could not write \(destination.lastPathComponent).")
        }
    }

    static func encrypt(_ source: URL, at destination: URL, password: String, ownerPassword: String?) throws {
        guard let document = PDFDocument(url: source) else {
            throw ConversionError.unreadable(source.lastPathComponent)
        }
        let options: [PDFDocumentWriteOption: Any] = [
            .userPasswordOption: password,
            .ownerPasswordOption: ownerPassword?.isEmpty == false ? ownerPassword! : password
        ]
        guard document.write(to: destination, withOptions: options) else {
            throw ConversionError.failed("Could not write the encrypted PDF.")
        }
    }

    /// Rewrites through an image pipeline, which is what "compress a PDF"
    /// means in practice: re-encode each page as a JPEG at lower quality.
    static func compress(
        _ source: URL,
        at destination: URL,
        quality: Double,
        cancellation: Cancellation? = nil,
        progress: ((Double) -> Void)? = nil
    ) throws {
        guard let document = PDFDocument(url: source) else {
            throw ConversionError.unreadable(source.lastPathComponent)
        }
        guard let firstPage = document.page(at: 0) else { return }
        var mediaBox = firstPage.bounds(for: .mediaBox)
        guard let context = CGContext(destination as CFURL, mediaBox: &mediaBox, nil) else {
            throw ConversionError.failed("Could not create the compressed PDF.")
        }
        for index in 0..<document.pageCount {
            try cancellation?.check()
            guard let page = document.page(at: index) else { continue }
            let bounds = page.bounds(for: .mediaBox)
            let scale: CGFloat = 150 / 72
            let width = Int(bounds.width * scale)
            let height = Int(bounds.height * scale)
            var box = bounds
            context.beginPDFPage([kCGPDFContextMediaBox: CGRectToData(box)] as CFDictionary)
            if width > 0, height > 0,
               let bitmap = CGContext(
                data: nil,
                width: width,
                height: height,
                bitsPerComponent: 8,
                bytesPerRow: 0,
                space: CGColorSpace(name: CGColorSpace.sRGB) ?? CGColorSpaceCreateDeviceRGB(),
                bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue
               ) {
                bitmap.setFillColor(CGColor(gray: 1, alpha: 1))
                bitmap.fill(CGRect(x: 0, y: 0, width: width, height: height))
                bitmap.scaleBy(x: scale, y: scale)
                bitmap.translateBy(x: -bounds.origin.x, y: -bounds.origin.y)
                page.draw(with: .mediaBox, to: bitmap)
                if let image = bitmap.makeImage(),
                   let jpeg = JPEGEncoder.encode(image, quality: quality) {
                    box = bounds
                    context.draw(jpeg, in: bounds)
                }
            }
            context.endPDFPage()
            progress?(Double(index + 1) / Double(document.pageCount))
        }
        context.closePDF()
    }

    /// Lays images out in a grid on a single page (Tangerine's "collage").
    static func collage(
        _ entries: [(image: CGImage, name: String)],
        at destination: URL,
        columns: Int,
        page: String,
        margin: CGFloat,
        captions: Bool,
        cancellation: Cancellation? = nil
    ) throws {
        guard !entries.isEmpty else { throw ConversionError.failed("Nothing to lay out.") }
        let pageSize: CGSize
        switch page {
        case "Letter": pageSize = CGSize(width: 612, height: 792)
        case "Square": pageSize = CGSize(width: 800, height: 800)
        default: pageSize = CGSize(width: 595, height: 842)
        }
        var mediaBox = CGRect(origin: .zero, size: pageSize)
        guard let context = CGContext(destination as CFURL, mediaBox: &mediaBox, nil) else {
            throw ConversionError.failed("Could not create \(destination.lastPathComponent).")
        }
        context.beginPDFPage([kCGPDFContextMediaBox: CGRectToData(mediaBox)] as CFDictionary)

        let columnCount = max(1, columns)
        let rowCount = Int(ceil(Double(entries.count) / Double(columnCount)))
        let usableWidth = pageSize.width - margin * 2
        let usableHeight = pageSize.height - margin * 2
        let cellWidth = usableWidth / CGFloat(columnCount)
        let cellHeight = usableHeight / CGFloat(max(1, rowCount))
        let captionHeight: CGFloat = captions ? 14 : 0

        for (index, entry) in entries.enumerated() {
            try cancellation?.check()
            let column = index % columnCount
            let row = index / columnCount
            let cell = CGRect(
                x: margin + CGFloat(column) * cellWidth,
                y: pageSize.height - margin - CGFloat(row + 1) * cellHeight,
                width: cellWidth,
                height: cellHeight
            ).insetBy(dx: 6, dy: 6)

            let imageArea = cell.insetBy(dx: 0, dy: captionHeight)
            let scale = min(imageArea.width / CGFloat(entry.image.width), imageArea.height / CGFloat(entry.image.height))
            let size = CGSize(width: CGFloat(entry.image.width) * scale, height: CGFloat(entry.image.height) * scale)
            context.draw(entry.image, in: CGRect(
                x: imageArea.midX - size.width / 2,
                y: imageArea.midY - size.height / 2,
                width: size.width,
                height: size.height
            ))

            if captions, let escaped = entry.name as CFString? {
                let font = CTFontCreateWithName("Helvetica" as CFString, 9, nil)
                let attributes: [NSAttributedString.Key: Any] = [
                    .font: font,
                    .foregroundColor: CGColor(gray: 0.35, alpha: 1)
                ]
                let line = CTLineCreateWithAttributedString(
                    NSAttributedString(string: escaped as String, attributes: attributes)
                )
                context.textPosition = CGPoint(x: cell.minX, y: cell.minY)
                CTLineDraw(line, context)
            }
        }
        context.endPDFPage()
        context.closePDF()
    }

    /// Rasterises a PDF page range to a single contact sheet, used by the
    /// "Info" summary.
    static func pageSize(_ source: URL) -> CGSize? {
        PDFDocument(url: source)?.page(at: 0)?.bounds(for: .mediaBox).size
    }

    static func pageCount(_ source: URL) -> Int {
        PDFDocument(url: source)?.pageCount ?? 0
    }
}

/// Small helper so the PDF compressor can produce JPEG bytes in memory.
enum JPEGEncoder {
    static func encode(_ image: CGImage, quality: Double) -> CGImage? {
        let flattened = Raster.flattened(image)
        guard let data = CFDataCreateMutable(nil, 0),
              let destination = CGImageDestinationCreateWithData(data, "public.jpeg" as CFString, 1, nil) else {
            return nil
        }
        CGImageDestinationAddImage(
            destination,
            flattened,
            [kCGImageDestinationLossyCompressionQuality: quality] as CFDictionary
        )
        guard CGImageDestinationFinalize(destination) else { return nil }
        guard let source = CGImageSourceCreateWithData(data, nil) else { return nil }
        return CGImageSourceCreateImageAtIndex(source, 0, nil)
    }
}

extension String {
    func leftPadded(to width: Int, with character: Character = "0") -> String {
        count >= width ? self : String(repeating: String(character), count: width - count) + self
    }
}