eric / piclead-client

import UIKit
import Vision
import CoreImage

/// Apple-native replacement for `OpenCvScanEngine`:
/// `VNDetectRectanglesRequest` finds the card quad, `CIPerspectiveCorrection` warps it.
/// Mirrors the Kotlin behavior: manual corners win (confidence 1), detected quad
/// scores 0.8, and when nothing is found the near-full-frame fallback (2% margin)
/// is warped with confidence 0.4 instead of failing.
final class VisionScanEngine: ScanEngine {

    func scan(_ source: UIImage, manualCorners: ScanCorners?) async throws -> ScanResult {
        guard let cgImage = source.cgImage ?? Self.renderCGImage(source) else {
            throw ScanException("Image invalide — impossible de décoder la photo")
        }
        let width = CGFloat(cgImage.width)
        let height = CGFloat(cgImage.height)

        let corners: ScanCorners
        let confidence: Float
        if let manualCorners {
            corners = manualCorners
            confidence = 1
        } else if let detected = await detectCardCorners(cgImage: cgImage, width: width, height: height) {
            corners = detected
            confidence = 0.8
        } else {
            corners = Self.fullFrameCorners(width: width, height: height)
            confidence = 0.4
        }

        let warped = try Self.perspectiveCorrect(cgImage: cgImage, corners: corners, imageHeight: height)
        return ScanResult(
            bitmap: warped,
            angleDegrees: Self.estimateSkewAngle(corners),
            confidence: confidence,
            corners: corners
        )
    }

    /// Returns detected corners, or nil when no plausible quad was found
    /// (Kotlin `detectDocumentCorners` null -> fallback path, never fatal).
    private func detectCardCorners(
        cgImage: CGImage,
        width: CGFloat,
        height: CGFloat
    ) async -> ScanCorners? {
        await withCheckedContinuation { continuation in
            DispatchQueue.global(qos: .userInitiated).async {
                let request = VNDetectRectanglesRequest()
                // Business cards: ~0.55 h/w landscape, up to square-ish when portrait.
                request.minimumAspectRatio = 0.3
                request.maximumAspectRatio = 1.0
                request.minimumSize = 0.2
                request.minimumConfidence = 0.6
                request.maximumObservations = 5
                request.quadratureTolerance = 20
                let handler = VNImageRequestHandler(cgImage: cgImage, orientation: .up)
                do {
                    try handler.perform([request])
                } catch {
                    continuation.resume(returning: nil)
                    return
                }
                guard let best = (request.results ?? []).max(by: { $0.confidence < $1.confidence }) else {
                    continuation.resume(returning: nil)
                    return
                }
                // Vision points are normalized, origin bottom-left -> pixels, origin top-left.
                func toPixels(_ p: CGPoint) -> CGPoint {
                    CGPoint(x: p.x * width, y: (1 - p.y) * height)
                }
                continuation.resume(returning: ScanCorners(
                    topLeft: toPixels(best.topLeft),
                    topRight: toPixels(best.topRight),
                    bottomRight: toPixels(best.bottomRight),
                    bottomLeft: toPixels(best.bottomLeft)
                ))
            }
        }
    }

    /// Kotlin `fullFrameCorners(marginRatio = 0.02)`.
    static func fullFrameCorners(
        width: CGFloat,
        height: CGFloat,
        marginRatio: CGFloat = 0.02
    ) -> ScanCorners {
        let mx = width * marginRatio
        let my = height * marginRatio
        let w = width - 1
        let h = height - 1
        return ScanCorners(
            topLeft: CGPoint(x: mx, y: my),
            topRight: CGPoint(x: w - mx, y: my),
            bottomRight: CGPoint(x: w - mx, y: h - my),
            bottomLeft: CGPoint(x: mx, y: h - my)
        )
    }

    /// Kotlin `perspectiveCorrect` equivalent via CIPerspectiveCorrection.
    static func perspectiveCorrect(
        cgImage: CGImage,
        corners: ScanCorners,
        imageHeight: CGFloat
    ) throws -> UIImage {
        // ScanCorners use a top-left origin; Core Image uses bottom-left.
        func ciPoint(_ p: CGPoint) -> CIVector {
            CIVector(x: p.x, y: imageHeight - p.y)
        }
        guard let filter = CIFilter(name: "CIPerspectiveCorrection") else {
            throw ScanException("CIPerspectiveCorrection indisponible")
        }
        filter.setValue(CIImage(cgImage: cgImage), forKey: kCIInputImageKey)
        filter.setValue(ciPoint(corners.topLeft), forKey: "inputTopLeft")
        filter.setValue(ciPoint(corners.topRight), forKey: "inputTopRight")
        filter.setValue(ciPoint(corners.bottomRight), forKey: "inputBottomRight")
        filter.setValue(ciPoint(corners.bottomLeft), forKey: "inputBottomLeft")
        guard let output = filter.outputImage,
              let corrected = CIContext().createCGImage(output, from: output.extent) else {
            throw ScanException("Correction de perspective impossible")
        }
        return UIImage(cgImage: corrected)
    }

    /// Kotlin `estimateSkewAngle`: angle of the top edge, top-left-origin coordinates.
    static func estimateSkewAngle(_ corners: ScanCorners) -> Float {
        let dx = corners.topRight.x - corners.topLeft.x
        let dy = corners.topRight.y - corners.topLeft.y
        return Float(atan2(dy, dx) * 180 / .pi)
    }

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