// chirami-ocr for macOS: reads Japanese in screenshots with Live Text
// (VisionKit's ImageAnalyzer), which unlike Vision's text request also reads
// vertical text (macOS 14+). ImageAnalyzer is Swift-only, hence this helper.
//
// Protocol (see internal/ocr/ocr.go): each request on stdin is a 4-byte
// big-endian length and that many bytes of PNG; each reply is one line of
// JSON on stdout. The helper exits when stdin closes.
import AppKit
import Vision
import VisionKit

struct Line: Encodable {
    let text: String
    let vertical: Bool?
}

struct Reply: Encodable {
    var lines: [Line] = []
    var engine = "livetext"
    var error: String?
}

func readExactly(_ n: Int) -> Data? {
    var data = Data()
    while data.count < n {
        let chunk = FileHandle.standardInput.readData(ofLength: n - data.count)
        if chunk.isEmpty { return nil }
        data.append(chunk)
    }
    return data
}

func isKana(_ s: String) -> Bool {
    !s.isEmpty && s.unicodeScalars.allSatisfy { (0x3040...0x30FF).contains($0.value) }
}

// furigana finds ruby lines: kana-only lines Vision sees directly above a
// line at least ~1.5x taller that they overlap horizontally. Live Text gives
// no geometry, so its transcript lines are matched against these by text.
// Vision cannot read vertical text, so ruby beside vertical columns stays.
func furigana(_ image: CGImage) -> Set<String> {
    let req = VNRecognizeTextRequest()
    req.recognitionLanguages = ["ja-JP"]
    req.recognitionLevel = .accurate
    try? VNImageRequestHandler(cgImage: image).perform([req])
    let obs = (req.results ?? []).compactMap { o -> (String, CGRect)? in
        guard let t = o.topCandidates(1).first?.string else { return nil }
        return (t, o.boundingBox) // normalised, origin bottom-left
    }
    var ruby = Set<String>()
    for (t, r) in obs where isKana(t) {
        for (_, b) in obs where b.height >= r.height * 1.5 {
            let overlaps = r.minX < b.maxX && b.minX < r.maxX
            let gap = r.minY - b.maxY
            if overlaps && gap > -r.height * 0.5 && gap < r.height * 1.5 {
                ruby.insert(t)
                break
            }
        }
    }
    return ruby
}

@main
struct Main {
    static func main() async {
        let analyzer = ImageAnalyzer()
        var config = ImageAnalyzer.Configuration([.text])
        config.locales = ["ja-JP", "en-US"]
        let out = FileHandle.standardOutput
        let enc = JSONEncoder()
        while let hdr = readExactly(4) {
            let n = hdr.reduce(0) { $0 << 8 | Int($1) }
            guard let png = readExactly(n) else { break }
            var reply = Reply()
            if let img = NSImage(data: png),
               let cg = img.cgImage(forProposedRect: nil, context: nil, hints: nil) {
                do {
                    let a = try await analyzer.analyze(img, orientation: .up, configuration: config)
                    let ruby = furigana(cg)
                    reply.lines = a.transcript.split(separator: "\n").map(String.init)
                        .filter { !ruby.contains($0) }
                        .map { Line(text: $0, vertical: nil) }
                } catch {
                    reply.error = "\(error)"
                }
            } else {
                reply.error = "not an image"
            }
            var b = (try? enc.encode(reply)) ?? Data(#"{"error":"encode"}"#.utf8)
            b.append(0x0A)
            out.write(b)
        }
    }
}