import AppKit
import Foundation
import Vision

struct OCRBox: Codable {
    let text: String
    let confidence: Float
    let x: Double
    let y: Double
    let width: Double
    let height: Double
}

func fail(_ message: String) -> Never {
    FileHandle.standardError.write(Data((message + "\n").utf8))
    exit(2)
}

guard CommandLine.arguments.count == 2 else {
    fail("usage: vision_ocr <image.png>")
}

guard let image = NSImage(contentsOfFile: CommandLine.arguments[1]) else {
    fail("could not load image")
}

var proposed = NSRect(origin: .zero, size: image.size)
guard let cgImage = image.cgImage(forProposedRect: &proposed, context: nil, hints: nil) else {
    fail("could not create CGImage")
}

let request = VNRecognizeTextRequest()
request.recognitionLevel = .accurate
request.recognitionLanguages = ["zh-Hant", "en-US"]
request.usesLanguageCorrection = true
request.minimumTextHeight = 0.012

do {
    try VNImageRequestHandler(cgImage: cgImage, orientation: .up).perform([request])
} catch {
    fail("Vision OCR failed: \(error)")
}

let imageWidth = Double(cgImage.width)
let imageHeight = Double(cgImage.height)
let boxes: [OCRBox] = (request.results ?? []).compactMap { observation in
    guard let candidate = observation.topCandidates(1).first else { return nil }
    let box = observation.boundingBox
    return OCRBox(
        text: candidate.string,
        confidence: candidate.confidence,
        x: box.minX * imageWidth,
        y: (1.0 - box.maxY) * imageHeight,
        width: box.width * imageWidth,
        height: box.height * imageHeight
    )
}

do {
    let data = try JSONEncoder().encode(boxes)
    FileHandle.standardOutput.write(data)
    FileHandle.standardOutput.write(Data("\n".utf8))
} catch {
    fail("could not encode OCR results: \(error)")
}
