eric / piclead-client

import UIKit
import ImageIO

/// Port of android `scan/ImageOrientation.kt`.
/// Normalizes a photo's orientation before scan/OCR:
/// - EXIF rotations (90/180/270) and mirrored variants
/// - camera rotation fallback when EXIF is absent/normal
/// - EXIF and camera rotation are never applied together (no double rotation)
enum ImageOrientation {

    struct OrientedBitmap {
        let bitmap: UIImage
        /// Applied quarter-turn degrees (0/90/180/270), mirrors excluded.
        let appliedRotationDegrees: Int
        /// true when the correction came from EXIF, false for the camera fallback.
        let fromExif: Bool
        let exifOrientation: CGImagePropertyOrientation
    }

    /// `UIImage(data:)` maps the JPEG EXIF orientation onto `imageOrientation`.
    static func fromJpegBytes(_ data: Data, cameraRotationDegrees: Int = 0) -> OrientedBitmap? {
        guard let image = UIImage(data: data) else { return nil }
        return orient(image, cameraRotationDegrees: cameraRotationDegrees)
    }

    static func fromFile(_ url: URL, cameraRotationDegrees: Int = 0) -> OrientedBitmap? {
        guard let data = try? Data(contentsOf: url) else { return nil }
        return fromJpegBytes(data, cameraRotationDegrees: cameraRotationDegrees)
    }

    /// Applies EXIF when present and not normal, otherwise the camera rotation.
    static func orient(_ image: UIImage, cameraRotationDegrees: Int = 0) -> OrientedBitmap {
        let exifOrientation = cgOrientation(from: image.imageOrientation)
        if exifOrientation != .up {
            return OrientedBitmap(
                bitmap: bakeOrientation(image),
                appliedRotationDegrees: exifRotationDegrees(exifOrientation),
                fromExif: true,
                exifOrientation: exifOrientation
            )
        }
        let deg = normalizeDegrees(cameraRotationDegrees)
        return OrientedBitmap(
            bitmap: deg == 0 ? image : rotate(image, degrees: deg),
            appliedRotationDegrees: deg,
            fromExif: false,
            exifOrientation: exifOrientation
        )
    }

    /// Kotlin `exifRotationDegrees` — rotation component per EXIF value.
    static func exifRotationDegrees(_ orientation: CGImagePropertyOrientation) -> Int {
        switch orientation {
        case .right, .leftMirrored:      // EXIF 6 (rotate 90) / 5 (transpose)
            return 90
        case .down, .downMirrored:       // EXIF 3 (rotate 180) / 4 (flip vertical)
            return 180
        case .left, .rightMirrored:      // EXIF 8 (rotate 270) / 7 (transverse)
            return 270
        case .up, .upMirrored:
            return 0
        }
    }

    /// Kotlin `normalizeDegrees` — quantizes any angle to the nearest quarter turn.
    static func normalizeDegrees(_ degrees: Int) -> Int {
        var d = degrees % 360
        if d < 0 { d += 360 }
        switch d {
        case 45..<135: return 90
        case 135..<225: return 180
        case 225..<315: return 270
        default: return 0
        }
    }

    static func cgOrientation(from orientation: UIImage.Orientation) -> CGImagePropertyOrientation {
        switch orientation {
        case .up: return .up
        case .upMirrored: return .upMirrored
        case .down: return .down
        case .downMirrored: return .downMirrored
        case .left: return .left
        case .leftMirrored: return .leftMirrored
        case .right: return .right
        case .rightMirrored: return .rightMirrored
        @unknown default: return .up
        }
    }

    /// Kotlin `applyExifTransform` equivalent: bakes the orientation (rotation and
    /// mirror) into the pixels so downstream Vision/CI passes see upright data.
    static func bakeOrientation(_ image: UIImage) -> UIImage {
        guard image.imageOrientation != .up else { return image }
        let format = UIGraphicsImageRendererFormat.default()
        format.scale = image.scale
        return UIGraphicsImageRenderer(size: image.size, format: format).image { _ in
            image.draw(in: CGRect(origin: .zero, size: image.size))
        }
    }

    /// Clockwise quarter-turn rotation (Kotlin `rotate`).
    static func rotate(_ image: UIImage, degrees: Int) -> UIImage {
        let swaps = degrees == 90 || degrees == 270
        let size = swaps
            ? CGSize(width: image.size.height, height: image.size.width)
            : image.size
        let format = UIGraphicsImageRendererFormat.default()
        format.scale = image.scale
        return UIGraphicsImageRenderer(size: size, format: format).image { context in
            let cg = context.cgContext
            cg.translateBy(x: size.width / 2, y: size.height / 2)
            cg.rotate(by: CGFloat(degrees) * .pi / 180)
            image.draw(in: CGRect(
                x: -image.size.width / 2,
                y: -image.size.height / 2,
                width: image.size.width,
                height: image.size.height
            ))
        }
    }
}