import Foundation
import PDFKit
import Vision
import AppKit

// 用法: swiftc -O ocr_pdf.swift -o /tmp/ocr_pdf && /tmp/ocr_pdf <pdf路径> > 结果.txt
// 逐页输出 "=== PAGE n ===" + 从上到下、从左到右排序的 OCR 文本
func ocrPDF(_ path: String) {
    guard let doc = PDFDocument(url: URL(fileURLWithPath: path)) else {
        print("ERROR: cannot open PDF: \(path)")
        return
    }
    let pageCount = doc.pageCount
    print("TOTAL_PAGES=\(pageCount)")
    for i in 0..<pageCount {
        guard let page = doc.page(at: i) else { continue }
        // 每页渲染成高分辨率 CGImage（2000x3000 足够 OCR），免去 PDF->PNG 中间步骤
        let thumb = page.thumbnail(of: CGSize(width: 2000, height: 3000), for: .mediaBox)
        guard let cg = thumb.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
            print("=== PAGE \(i+1): render failed ===")
            continue
        }
        let request = VNRecognizeTextRequest()
        request.recognitionLevel = .accurate
        request.recognitionLanguages = ["zh-Hans", "en-US"]
        request.usesLanguageCorrection = true
        let handler = VNImageRequestHandler(cgImage: cg, options: [:])
        do {
            try handler.perform([request])
        } catch {
            print("=== PAGE \(i+1): OCR error \(error) ===")
            continue
        }
        print("=== PAGE \(i+1) ===")
        guard let obs = request.results else { continue }
        // boundingBox 原点在左下角：midY 降序 = 从上到下；同行 minX 升序 = 从左到右
        let sorted = obs.sorted { a, b in
            if abs(a.boundingBox.midY - b.boundingBox.midY) > 0.012 {
                return a.boundingBox.midY > b.boundingBox.midY
            }
            return a.boundingBox.minX < b.boundingBox.minX
        }
        for o in sorted {
            if let c = o.topCandidates(1).first {
                print(c.string)
            }
        }
    }
}

if CommandLine.arguments.count > 1 {
    ocrPDF(CommandLine.arguments[1])
} else {
    print("Usage: swift ocr_pdf.swift <pdf>")
}