import AppKit
import Foundation
import Vision

struct OCRLine {
    let text: String
    let x: CGFloat
    let y: CGFloat
}

func pageNumber(from path: String) -> Int {
    let name = URL(fileURLWithPath: path).deletingPathExtension().lastPathComponent
    let digits = name.replacingOccurrences(of: "[^0-9]", with: "", options: .regularExpression)
    return Int(digits) ?? 0
}

func loadCGImage(_ path: String) -> CGImage? {
    let url = URL(fileURLWithPath: path)
    guard let image = NSImage(contentsOf: url),
          let tiff = image.tiffRepresentation,
          let bitmap = NSBitmapImageRep(data: tiff) else {
        return nil
    }
    return bitmap.cgImage
}

func recognize(_ path: String) throws -> [OCRLine] {
    guard let cgImage = loadCGImage(path) else {
        throw NSError(domain: "VisionOCR", code: 1, userInfo: [NSLocalizedDescriptionKey: "Could not load image: \(path)"])
    }

    var lines: [OCRLine] = []
    var requestError: Error?
    let request = VNRecognizeTextRequest { request, error in
        if let error = error {
            requestError = error
            return
        }
        let observations = (request.results as? [VNRecognizedTextObservation]) ?? []
        lines = observations.compactMap { observation in
            guard let candidate = observation.topCandidates(1).first else { return nil }
            let text = candidate.string.trimmingCharacters(in: .whitespacesAndNewlines)
            guard !text.isEmpty else { return nil }
            return OCRLine(text: text, x: observation.boundingBox.minX, y: observation.boundingBox.maxY)
        }
    }

    request.recognitionLevel = .accurate
    request.usesLanguageCorrection = true
    request.recognitionLanguages = ["zh-Hans", "en-US"]

    let handler = VNImageRequestHandler(cgImage: cgImage, options: [:])
    try handler.perform([request])
    if let requestError = requestError {
        throw requestError
    }

    return lines.sorted {
        if abs($0.y - $1.y) > 0.012 {
            return $0.y > $1.y
        }
        return $0.x < $1.x
    }
}

let args = Array(CommandLine.arguments.dropFirst())
guard args.count >= 2 else {
    FileHandle.standardError.write(Data("Usage: vision_ocr.swift OUTPUT.md IMAGE...\n".utf8))
    exit(2)
}

let output = args[0]
let imagePaths = args.dropFirst().sorted {
    let lhs = pageNumber(from: $0)
    let rhs = pageNumber(from: $1)
    return lhs == rhs ? $0 < $1 : lhs < rhs
}

var result = "# GB/T 33923-2017 OCR 文本版\n\n"
result += "> 说明：此文本由 PDF 页面图像 OCR 生成，用于阅读、搜索和复制。公式、表格、上下标和个别专业符号可能需要对照原 PDF 校核。\n\n"

for (index, path) in imagePaths.enumerated() {
    let page = pageNumber(from: path)
    FileHandle.standardError.write(Data("OCR page \(index + 1)/\(imagePaths.count): \(URL(fileURLWithPath: path).lastPathComponent)\n".utf8))
    do {
        let lines = try recognize(path)
        result += "\n\n## 第 \(page) 页\n\n"
        result += lines.map(\.text).joined(separator: "\n")
        result += "\n"
    } catch {
        result += "\n\n## 第 \(page) 页\n\n[OCR 失败：\(error.localizedDescription)]\n"
    }
}

try result.write(toFile: output, atomically: true, encoding: .utf8)