shotdeck/Sources/ShotdeckCore/PDF/PDFComposer.swift

import AppKit
import CoreGraphics
import CoreText
import Foundation
import ImageIO
import os

private let pdfLog = Logger(subsystem: "ai.flowmaster.shotdeck", category: "PDF")

public struct PDFComposer: Sendable {
    public init() {}

    @discardableResult
    public func compose(
        session: CaptureSession,
        imageURL: (Capture) -> URL,
        title: String,
        to outputURL: URL
    ) throws -> Int {
        if session.isEmpty {
            throw ShotdeckError.pdfCompositionFailed(reason: "session has no captures")
        }

        var loaded: [(Capture, CGImage)] = []
        loaded.reserveCapacity(session.captures.count)
        for capture in session.captures {
            let url = imageURL(capture)
            if let image = Self.loadCGImage(from: url) {
                loaded.append((capture, image))
            } else {
                pdfLog.error(
                    "skipping capture \(capture.id.uuidString, privacy: .public) seq \(capture.sequence, privacy: .public): image unreadable at \(url.path, privacy: .public)"
                )
            }
        }

        guard !loaded.isEmpty else {
            throw ShotdeckError.pdfCompositionFailed(reason: "no readable images in session")
        }

        let pageCount = loaded.count
        let tmpURL = outputURL
            .deletingLastPathComponent()
            .appendingPathComponent(".tmp-\(UUID().uuidString)-\(outputURL.lastPathComponent)")

        do {
            try Self.writePDF(
                loaded: loaded,
                pageCount: pageCount,
                title: title,
                sessionID: session.id,
                to: tmpURL
            )

            let handle = try FileHandle(forWritingTo: tmpURL)
            try handle.synchronize()
            try handle.close()

            if FileManager.default.fileExists(atPath: outputURL.path) {
                try FileManager.default.removeItem(at: outputURL)
            }
            try FileManager.default.moveItem(at: tmpURL, to: outputURL)
            return pageCount
        } catch {
            try? FileManager.default.removeItem(at: tmpURL)
            if let shotdeck = error as? ShotdeckError {
                throw shotdeck
            }
            throw ShotdeckError.pdfCompositionFailed(reason: error.localizedDescription)
        }
    }

    public static func fileName(for session: CaptureSession) -> String {
        "Shotdeck-\(DubaiTime.fileStamp(session.createdAt)).pdf"
    }

    private static func writePDF(
        loaded: [(Capture, CGImage)],
        pageCount: Int,
        title: String,
        sessionID: UUID,
        to tmpURL: URL
    ) throws {
        guard let consumer = CGDataConsumer(url: tmpURL as CFURL) else {
            throw ShotdeckError.pdfCompositionFailed(reason: "cannot open output location")
        }

        let auxiliaryInfo: [String: Any] = [
            kCGPDFContextCreator as String: "Shotdeck",
            kCGPDFContextTitle as String: title,
            kCGPDFContextSubject as String: sessionID.uuidString.lowercased(),
        ]

        guard let context = CGContext(
            consumer: consumer,
            mediaBox: nil,
            auxiliaryInfo as CFDictionary
        ) else {
            throw ShotdeckError.pdfCompositionFailed(reason: "cannot open output location")
        }

        let lightGrey = CGColor(gray: 0.75, alpha: 1)
        let black = CGColor(gray: 0, alpha: 1)
        let timestampColor = CGColor(gray: 0.45, alpha: 1)

        for pageIndex in 1...pageCount {
            let (capture, cgImage) = loaded[pageIndex - 1]
            let layout = PageLayout(capture: capture, pageIndex: pageIndex, pageCount: pageCount)

            context.beginPDFPage(Self.pageInfo(mediaBox: layout.pageRect))

            context.draw(cgImage, in: layout.imageRect)

            context.setStrokeColor(lightGrey)
            context.setLineWidth(0.5)
            context.move(to: CGPoint(x: layout.headerRect.minX, y: layout.headerRect.minY))
            context.addLine(to: CGPoint(x: layout.headerRect.maxX, y: layout.headerRect.minY))
            context.strokePath()

            drawText(
                layout.pageNumberText,
                font: layout.pageNumberFont,
                color: black,
                at: layout.pageNumberOrigin,
                in: context
            )
            drawText(
                layout.timestampText,
                font: layout.timestampFont,
                color: timestampColor,
                at: layout.timestampOrigin,
                in: context
            )
            drawText(
                "PASS",
                font: layout.tickLabelFont,
                color: black,
                at: layout.passLabelOrigin,
                in: context
            )
            drawText(
                "FAIL",
                font: layout.tickLabelFont,
                color: black,
                at: layout.failLabelOrigin,
                in: context
            )

            strokeTickBox(layout.passTickBox, in: context, color: black)
            strokeTickBox(layout.failTickBox, in: context, color: black)

            context.endPDFPage()
        }

        context.closePDF()
    }

    private static func pageInfo(mediaBox: CGRect) -> CFDictionary {
        var box = mediaBox
        let data = Data(bytes: &box, count: MemoryLayout<CGRect>.size)
        return [kCGPDFContextMediaBox as String: data] as CFDictionary
    }

    private static func drawText(
        _ string: String,
        font: NSFont,
        color: CGColor,
        at origin: CGPoint,
        in context: CGContext
    ) {
        let attributes: [CFString: Any] = [
            kCTFontAttributeName: font,
            kCTForegroundColorAttributeName: color,
        ]
        guard let attributed = CFAttributedStringCreate(
            nil,
            string as CFString,
            attributes as CFDictionary
        ) else {
            return
        }
        let line = CTLineCreateWithAttributedString(attributed)
        context.textMatrix = .identity
        context.textPosition = origin
        CTLineDraw(line, context)
    }

    private static func strokeTickBox(_ box: CGRect, in context: CGContext, color: CGColor) {
        let path = CGPath(
            roundedRect: box,
            cornerWidth: 2,
            cornerHeight: 2,
            transform: nil
        )
        context.addPath(path)
        context.setStrokeColor(color)
        context.setLineWidth(1.0)
        context.strokePath()
    }

    private static func loadCGImage(from url: URL) -> CGImage? {
        guard let source = CGImageSourceCreateWithURL(url as CFURL, nil),
              CGImageSourceGetCount(source) > 0
        else {
            return nil
        }
        return CGImageSourceCreateImageAtIndex(source, 0, nil)
    }
}