// radio_bug.swift
// Reproduce PDFKit radio button /AS bug
// Run with: swift radio_bug.swift
//
// Expected: The annotation with buttonWidgetState = .onState should have
//           its /AS set to the corresponding buttonWidgetStateString in the PDF.
// Actual:   Only the FIRST annotation added via page.addAnnotation() gets a
//           correct /AS — all others always get /AS /Off.
//
// Root cause (confirmed via method swizzling):
//   The bug is entirely in dataRepresentation()'s serialization path.
//   Setting buttonWidgetState = .onState before calling dataRepresentation()
//   works correctly — the in-memory state is right.
//   However, dataRepresentation() internally calls setButtonWidgetState:.onState
//   on each annotation in turn to generate its appearance streams, which
//   triggers PDFKit's radio-group exclusivity logic and silently resets all
//   other annotations in the group to .offState.  By the time each annotation's
//   /AS is written, its state has already been clobbered by the appearance-
//   generation pass of a preceding annotation.
//
//   As a further consequence, dataRepresentation() corrupts the caller's own
//   in-memory PDFAnnotation objects as a side effect (see test at the bottom).

import PDFKit
import Foundation

func makeRadioButton(optionId: String, isSelected: Bool) -> PDFAnnotation {
    let ann = PDFAnnotation(
        bounds: CGRect(x: 0, y: 0, width: 20, height: 20),
        forType: .widget,
        withProperties: nil
    )
    ann.widgetFieldType = .button
    ann.widgetControlType = .radioButtonControl
    ann.fieldName = "Choice"
    ann.buttonWidgetStateString = optionId
    ann.buttonWidgetState = isSelected ? .onState : .offState
    return ann
}

var allPassed = true

let outputDir = URL(fileURLWithPath: FileManager.default.currentDirectoryPath)

for selectedOption in ["A", "B", "C"] {
    let pdf = PDFDocument()
    let page = PDFPage()
    pdf.insert(page, at: 0)

    // Add labels and radio buttons for visual inspection
    let labelTitle = PDFAnnotation(bounds: CGRect(x: 30, y: 720, width: 200, height: 20), forType: .freeText, withProperties: nil)
    labelTitle.contents = "Select \(selectedOption) — Expected: \(selectedOption) selected"
    labelTitle.font = NSFont.systemFont(ofSize: 12)
    labelTitle.fontColor = .black
    labelTitle.color = .clear
    page.addAnnotation(labelTitle)

    for (index, optionId) in ["A", "B", "C"].enumerated() {
        let y = 680 - index * 40

        // Radio button
        let radio = makeRadioButton(optionId: optionId, isSelected: selectedOption == optionId)
        radio.bounds = CGRect(x: 50, y: y, width: 20, height: 20)
        page.addAnnotation(radio)

        // Option label
        let label = PDFAnnotation(bounds: CGRect(x: 80, y: y, width: 300, height: 20), forType: .freeText, withProperties: nil)
        label.contents = "Option \(optionId)\(selectedOption == optionId ? "  ← should be selected" : "")"
        label.font = NSFont.systemFont(ofSize: 11)
        label.fontColor = selectedOption == optionId ? .red : .black
        label.color = .clear
        page.addAnnotation(label)
    }

    // Capture in-memory state BEFORE serialization
    let radioAnns = page.annotations.filter { $0.widgetControlType == .radioButtonControl }
    func stateStr(_ anns: [PDFAnnotation]) -> String {
        anns.map { "\($0.buttonWidgetStateString)=\($0.buttonWidgetState == .onState ? "ON" : "off")" }
            .joined(separator: ", ")
    }
    print("  buttonWidgetState before dataRepresentation(): \(stateStr(radioAnns))")

    print("   Calling dataRepresentation() to serialize PDF")
    guard let data = pdf.dataRepresentation() else { continue }

    // Capture in-memory state AFTER serialization — should be unchanged, but isn't
    print("  buttonWidgetState after  dataRepresentation(): \(stateStr(radioAnns))")

    // Save PDF to disk
    let pdfURL = outputDir.appendingPathComponent("\(selectedOption).pdf")
    try? data.write(to: pdfURL)
    print("  Written: \(pdfURL.path)")

    // Verify /AS values in raw PDF bytes
    guard let pdfStr = String(data: data, encoding: .isoLatin1) else { continue }
    let asValues = pdfStr
        .components(separatedBy: "/AS ")
        .dropFirst()
        .map { String($0.prefix(6).trimmingCharacters(in: .whitespaces)) }

    let pass = asValues.contains { $0.hasPrefix("/\(selectedOption)") }
    if !pass { allPassed = false }
    print("Select \(selectedOption): /AS values = \(asValues) → \(pass ? "PASS" : "FAIL (bug)")")
    print()
}

print(allPassed ? "All tests passed." : "Test failed.")

// --- Workaround demonstration: select C by adding C first ---
print()
print("--- Workaround: add selected annotation first ---")
do {
    let selectedOption = "C"
    let pdf = PDFDocument()
    let page = PDFPage()
    pdf.insert(page, at: 0)

    let titleLabel = PDFAnnotation(bounds: CGRect(x: 30, y: 720, width: 400, height: 20), forType: .freeText, withProperties: nil)
    titleLabel.contents = "Workaround: Select \(selectedOption) — add selected annotation first"
    titleLabel.font = NSFont.systemFont(ofSize: 12)
    titleLabel.fontColor = .black
    titleLabel.color = .clear
    page.addAnnotation(titleLabel)

    // Workaround: add the selected option (C) first, then the others
    let orderedOptions = ["C", "A", "B"]  // selected comes first
    let displayPositions: [String: Int] = ["A": 0, "B": 1, "C": 2]

    for optionId in orderedOptions {
        let index = displayPositions[optionId]!
        let y = 680 - index * 40

        let radio = makeRadioButton(optionId: optionId, isSelected: optionId == selectedOption)
        radio.bounds = CGRect(x: 50, y: y, width: 20, height: 20)
        page.addAnnotation(radio)

        let label = PDFAnnotation(bounds: CGRect(x: 80, y: y, width: 400, height: 20), forType: .freeText, withProperties: nil)
        label.contents = "Option \(optionId)\(optionId == selectedOption ? "  ← selected (added first as workaround)" : "")"
        label.font = NSFont.systemFont(ofSize: 11)
        label.fontColor = optionId == selectedOption ? .blue : .black
        label.color = .clear
        page.addAnnotation(label)
    }

    guard let data = pdf.dataRepresentation() else { exit(1) }

    let pdfURL = outputDir.appendingPathComponent("Workaround_C.pdf")
    try? data.write(to: pdfURL)
    print("Written: \(pdfURL.path)")

    guard let pdfStr = String(data: data, encoding: .isoLatin1) else { exit(1) }
    let asValues = pdfStr
        .components(separatedBy: "/AS ")
        .dropFirst()
        .map { String($0.prefix(6).trimmingCharacters(in: .whitespaces)) }

    let pass = asValues.contains { $0.hasPrefix("/\(selectedOption)") }
    print("Workaround Select \(selectedOption): /AS values = \(asValues) → \(pass ? "PASS (workaround works)" : "FAIL")")
}

exit(allPassed ? 0 : 1)