#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
import re
import subprocess
import sys
from pathlib import Path


SWIFT_CODE = r'''
import Foundation
import Vision
import AppKit

let path = CommandLine.arguments[1]
let url = URL(fileURLWithPath: path)
guard let img = NSImage(contentsOf: url),
      let tiff = img.tiffRepresentation,
      let bitmap = NSBitmapImageRep(data: tiff),
      let cg = bitmap.cgImage else {
    fputs("image load failed\n", stderr)
    exit(3)
}

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

var lines: [Line] = []
var requestError: Error?
let request = VNRecognizeTextRequest { req, err in
    if let err = err {
        requestError = err
        return
    }
    let observations = (req.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 Line(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: cg, options: [:])
do {
    try handler.perform([request])
} catch {
    fputs("\(error.localizedDescription)\n", stderr)
    exit(4)
}
if let requestError = requestError {
    fputs("\(requestError.localizedDescription)\n", stderr)
    exit(5)
}

let sorted = lines.sorted {
    if abs($0.y - $1.y) > 0.012 {
        return $0.y > $1.y
    }
    return $0.x < $1.x
}
for line in sorted {
    print(line.text)
}
'''


def page_number(path: Path) -> int:
    match = re.search(r"(\d+)", path.stem)
    return int(match.group(1)) if match else 0


def run_ocr(image: Path, module_cache: Path) -> tuple[str, str | None]:
    module_cache.mkdir(parents=True, exist_ok=True)
    env = os.environ.copy()
    env["CLANG_MODULE_CACHE_PATH"] = str(module_cache)
    proc = subprocess.run(
        ["swift", "-e", SWIFT_CODE, str(image)],
        text=True,
        capture_output=True,
        check=False,
        env=env,
    )
    if proc.returncode != 0:
        return "", proc.stderr.strip() or f"swift exited {proc.returncode}"
    return proc.stdout.strip(), None


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("output", type=Path)
    parser.add_argument("images", nargs="+", type=Path)
    args = parser.parse_args()

    images = sorted(args.images, key=lambda p: (page_number(p), str(p)))
    args.output.parent.mkdir(parents=True, exist_ok=True)
    module_cache = args.output.parent.parent / "tmp" / "swift_module_cache"

    with args.output.open("w", encoding="utf-8") as out:
        out.write("# GB/T 33923-2017 OCR 文本版\n\n")
        out.write("> 说明：此文本由 PDF 页面图像 OCR 生成，用于阅读、搜索和复制。公式、表格、上下标和个别专业符号可能需要对照原 PDF 校核。\n\n")
        for index, image in enumerate(images, 1):
            page = page_number(image)
            print(f"OCR page {index}/{len(images)}: {image.name}", file=sys.stderr, flush=True)
            text, error = run_ocr(image, module_cache)
            out.write(f"\n\n## 第 {page} 页\n\n")
            if error:
                out.write(f"[OCR 失败：{error}]\n")
            elif text:
                out.write(text)
                out.write("\n")
            else:
                out.write("[OCR 未识别出文本]\n")
            out.flush()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())