import AppKit
import Foundation
import Vision

struct OCRLine: Codable {
  let text: String
  let x: Double
  let y: Double
  let width: Double
  let height: Double
}

struct OCRPage: Codable {
  let grade: String
  let file: String
  let path: String
  let lines: [OCRLine]
}

let args = CommandLine.arguments
guard args.count == 3 else {
  fputs("usage: research_ocr_2026_fall <source-root> <output-json>\n", stderr)
  exit(2)
}

let sourceRoot = URL(fileURLWithPath: args[1], isDirectory: true)
let outputURL = URL(fileURLWithPath: args[2])
let manager = FileManager.default
let allowedExtensions = Set(["png", "jpg", "jpeg"])

let gradeURLs = try manager.contentsOfDirectory(
  at: sourceRoot,
  includingPropertiesForKeys: nil,
  options: [.skipsHiddenFiles]
).filter { $0.hasDirectoryPath }.sorted { $0.lastPathComponent < $1.lastPathComponent }

var pages: [OCRPage] = []
for gradeURL in gradeURLs {
  let imageURLs = try manager.contentsOfDirectory(
    at: gradeURL,
    includingPropertiesForKeys: nil,
    options: [.skipsHiddenFiles]
  ).filter { allowedExtensions.contains($0.pathExtension.lowercased()) }
    .sorted { $0.lastPathComponent < $1.lastPathComponent }

  for (index, imageURL) in imageURLs.enumerated() {
    guard let image = NSImage(contentsOf: imageURL),
          let cgImage = image.cgImage(forProposedRect: nil, context: nil, hints: nil) else {
      fputs("failed to load \(imageURL.path)\n", stderr)
      exit(1)
    }

    let request = VNRecognizeTextRequest()
    request.recognitionLevel = .accurate
    request.usesLanguageCorrection = true
    request.recognitionLanguages = ["zh-Hans", "en-US"]
    request.minimumTextHeight = 0.006
    try VNImageRequestHandler(cgImage: cgImage, options: [:]).perform([request])

    let lines = (request.results ?? []).compactMap { observation -> OCRLine? in
      guard let candidate = observation.topCandidates(1).first else { return nil }
      let box = observation.boundingBox
      return OCRLine(
        text: candidate.string,
        x: Double(box.minX),
        y: Double(box.minY),
        width: Double(box.width),
        height: Double(box.height)
      )
    }.sorted {
      let leftY = Int($0.y * 1000)
      let rightY = Int($1.y * 1000)
      if leftY != rightY { return leftY > rightY }
      return $0.x < $1.x
    }

    pages.append(OCRPage(
      grade: gradeURL.lastPathComponent,
      file: imageURL.lastPathComponent,
      path: imageURL.path,
      lines: lines
    ))
    fputs("[\(gradeURL.lastPathComponent)] \(index + 1)/\(imageURLs.count) \(imageURL.lastPathComponent)\n", stderr)
  }
}

let encoder = JSONEncoder()
encoder.outputFormatting = [.prettyPrinted, .sortedKeys, .withoutEscapingSlashes]
try encoder.encode(pages).write(to: outputURL, options: .atomic)
fputs("wrote \(pages.count) pages to \(outputURL.path)\n", stderr)