eric / piclead-client

import UIKit
import Vision
import CoreImage

/// Apple-native replacement for `TesseractOcrEngine` — 100% on-device, no network.
/// Mirrors the Kotlin contract: tries the four quarter-turn orientations
/// (`OcrOrientation.recognizeBest`) and post-processes raw text through
/// `OcrPostProcessor.process` before returning.
final class VisionOcrEngine: OcrEngine {

    /// Tesseract language code -> BCP-47 code for `VNRecognizeTextRequest`.
    /// Covers the app's shipped languages (`fra+deu+eng+spa+por+ita+pol`) plus `nld`.
    private static let tesseractToBcp47: [String: String] = [
        "fra": "fr-FR",
        "deu": "de-DE",
        "eng": "en-US",
        "spa": "es-ES",
        "por": "pt-BR",
        "ita": "it-IT",
        "nld": "nl-NL",
        "pol": "pl-PL",
    ]

    /// Mirrors `OcrOrientation.angles` = [0, 90, 180, 270] (clockwise).
    private static let orientations: [CGImagePropertyOrientation] = [.up, .right, .down, .left]

    private let recognitionLanguages: [String]

    /// `languages` keeps the Kotlin constructor contract ("fra+deu+...").
    /// Codes Vision cannot handle on this OS version are dropped at runtime.
    init(languages: String = "fra+deu+eng+spa+por+ita+pol") {
        let requested = languages
            .split(separator: "+")
            .compactMap { Self.tesseractToBcp47[String($0).lowercased()] }
        recognitionLanguages = Self.filterSupported(requested)
    }

    func recognize(_ image: UIImage) async throws -> OcrResult {
        guard let cgImage = image.cgImage ?? Self.renderCGImage(image) else {
            throw OcrException("Image invalide — impossible de décoder la capture")
        }
        var best: OcrResult?
        var bestScore = Int.min
        for orientation in Self.orientations {
            guard let result = try? await recognizeOnce(cgImage: cgImage, orientation: orientation) else {
                continue
            }
            let score = OcrOrientation.score(result)
            if score > bestScore {
                bestScore = score
                best = result
            }
            // Good enough: email + phone found (same early-exit as Kotlin).
            if !result.emails.isEmpty && !result.phones.isEmpty { break }
        }
        guard let best else { throw OcrException("OCR vide — reprenez la photo") }
        return best
    }

    /// Mirrors `TesseractOcrEngine.recognizeOnce`: raw text -> OcrPostProcessor,
    /// plus the spatial lines (boxes + confidence) for the layout-aware parser.
    private func recognizeOnce(
        cgImage: CGImage,
        orientation: CGImagePropertyOrientation
    ) async throws -> OcrResult {
        let spatial = try await performVision(cgImage: cgImage, orientation: orientation)
        let raw = spatial.map(\.text).joined(separator: "\n")
        let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
        if trimmed.isEmpty { throw OcrException("OCR vide") }
        var result = OcrPostProcessor.process(trimmed)
        result.spatialLines = spatial
        return result
    }

    private func performVision(
        cgImage: CGImage,
        orientation: CGImagePropertyOrientation
    ) async throws -> [OcrLine] {
        let languages = recognitionLanguages
        return try await withCheckedThrowingContinuation { continuation in
            DispatchQueue.global(qos: .userInitiated).async {
                let request = VNRecognizeTextRequest()
                request.recognitionLevel = .accurate
                request.usesLanguageCorrection = true
                if !languages.isEmpty {
                    request.recognitionLanguages = languages
                }
                let handler = VNImageRequestHandler(cgImage: cgImage, orientation: orientation)
                do {
                    try handler.perform([request])
                    continuation.resume(returning: Self.spatialLines(request.results ?? []))
                } catch {
                    continuation.resume(throwing: OcrException("Vision OCR a échoué", underlying: error))
                }
            }
        }
    }

    /// Observations → `OcrLine`s sorted top-to-bottom then left-to-right.
    ///
    /// Vision boxes are normalized (0-1) with a bottom-left origin; `OcrLine`
    /// expects a y-down frame where only RELATIVE positions matter, so the
    /// boxes are mapped into a virtual 1000×1000 space with a flipped y axis.
    /// Vision confidence (0-1) is scaled to the Tesseract range (0-100).
    private static func spatialLines(_ observations: [VNRecognizedTextObservation]) -> [OcrLine] {
        observations
            .compactMap { observation -> OcrLine? in
                guard let candidate = observation.topCandidates(1).first else { return nil }
                let text = candidate.string.trimmingCharacters(in: .whitespacesAndNewlines)
                guard !text.isEmpty else { return nil }
                let box = observation.boundingBox
                return OcrLine(
                    text: text,
                    box: OcrBox(
                        left: Int(box.minX * 1000),
                        top: Int((1 - box.maxY) * 1000),
                        right: Int(box.maxX * 1000),
                        bottom: Int((1 - box.minY) * 1000)
                    ),
                    confidence: candidate.confidence * 100
                )
            }
            .sorted { a, b in
                if abs(a.box.top - b.box.top) > 10 { return a.box.top < b.box.top }
                return a.box.left < b.box.left
            }
    }

    /// Keeps only languages Vision supports on this OS (exact code, then base-language match).
    private static func filterSupported(_ requested: [String]) -> [String] {
        let probe = VNRecognizeTextRequest()
        probe.recognitionLevel = .accurate
        guard let supported = try? probe.supportedRecognitionLanguages() else { return requested }
        let exact = Set(supported.map { $0.lowercased() })
        let bases = Set(supported.map { $0.split(separator: "-").first.map(String.init)?.lowercased() ?? "" })
        return requested.filter { code in
            exact.contains(code.lowercased())
                || bases.contains(code.split(separator: "-").first.map(String.init)?.lowercased() ?? "")
        }
    }

    /// Renders CIImage-backed UIImages to a CGImage.
    private static func renderCGImage(_ image: UIImage) -> CGImage? {
        guard let ciImage = image.ciImage else { return nil }
        return CIContext().createCGImage(ciImage, from: ciImage.extent)
    }
}