PlantGuide/PlantGuide/ML/Preprocessing/ImagePreprocessor.swift

//
//  ImagePreprocessor.swift
//  PlantGuide
//
//  Created by Trey Tartt on 1/21/26.
//

import UIKit
import CoreGraphics
import ImageIO

// MARK: - Image Preprocessor Protocol

/// Protocol for preprocessing images before classification
protocol ImagePreprocessorProtocol: Sendable {
    /// Preprocesses a UIImage for model input
    /// - Parameter image: The source UIImage to preprocess
    /// - Returns: A CGImage ready for classification
    /// - Throws: If preprocessing fails
    func preprocess(_ image: UIImage) async throws -> CGImage
}

// MARK: - Image Preprocessor Error

enum ImagePreprocessorError: LocalizedError, Equatable {
    case invalidImage
    case corruptData
    case unsupportedFormat
    case dimensionsTooSmall(width: Int, height: Int)
    case colorSpaceConversionFailed
    case cgImageCreationFailed

    var errorDescription: String? {
        switch self {
        case .invalidImage:
            return "The provided image is invalid or nil."
        case .corruptData:
            return "The image data is corrupt or unreadable."
        case .unsupportedFormat:
            return "The image format is not supported."
        case .dimensionsTooSmall(let width, let height):
            return "Image dimensions (\(width)x\(height)) are below the minimum required (224x224)."
        case .colorSpaceConversionFailed:
            return "Failed to convert image to sRGB color space."
        case .cgImageCreationFailed:
            return "Failed to create CGImage from the provided image."
        }
    }
}

// MARK: - Image Preprocessor Configuration

/// Configuration for image preprocessing operations.
struct ImagePreprocessorConfiguration: Sendable {
    /// Minimum width required for ML model input.
    let minimumWidth: Int

    /// Minimum height required for ML model input.
    let minimumHeight: Int

    /// Whether to enforce sRGB color space conversion.
    let requiresSRGB: Bool

    /// Default configuration for plant identification model (224x224 RGB input).
    static let `default` = ImagePreprocessorConfiguration(
        minimumWidth: 224,
        minimumHeight: 224,
        requiresSRGB: true
    )

    init(minimumWidth: Int = 224, minimumHeight: Int = 224, requiresSRGB: Bool = true) {
        self.minimumWidth = minimumWidth
        self.minimumHeight = minimumHeight
        self.requiresSRGB = requiresSRGB
    }
}

// MARK: - Image Preprocessor

/// Prepares images for Core ML model inference.
///
/// This struct handles:
/// - EXIF orientation correction
/// - Color space conversion to sRGB
/// - Dimension validation (minimum 224x224)
/// - Various input formats (UIImage, Data)
///
/// The Vision framework handles actual resizing to model input dimensions.
struct ImagePreprocessor: ImagePreprocessorProtocol, Sendable {
    // MARK: - Properties

    private let configuration: ImagePreprocessorConfiguration

    // MARK: - Initialization

    init(configuration: ImagePreprocessorConfiguration = .default) {
        self.configuration = configuration
    }

    // MARK: - ImagePreprocessorProtocol Conformance

    /// Preprocesses a UIImage for model input (async throwing variant for protocol conformance).
    /// - Parameter image: The source UIImage to preprocess.
    /// - Returns: A CGImage ready for classification.
    /// - Throws: `ImagePreprocessorError` if preprocessing fails.
    func preprocess(_ image: UIImage) async throws -> CGImage {
        try prepareWithValidation(image: image)
    }

    // MARK: - Convenience Methods

    /// Prepares a UIImage for ML model input.
    /// - Parameter image: The source UIImage to prepare.
    /// - Returns: A CGImage ready for Vision framework processing, or nil if preparation fails.
    func prepare(image: UIImage) -> CGImage? {
        // Handle nil cgImage case
        guard let cgImage = extractCGImage(from: image) else {
            return nil
        }

        return processImage(cgImage)
    }

    /// Prepares image data for ML model input.
    /// - Parameter data: Raw image data (JPEG, PNG, HEIC, etc.).
    /// - Returns: A CGImage ready for Vision framework processing, or nil if preparation fails.
    func prepare(data: Data) -> CGImage? {
        // Validate data is not empty
        guard !data.isEmpty else {
            return nil
        }

        // Create image source from data
        guard let imageSource = CGImageSourceCreateWithData(data as CFData, nil) else {
            return nil
        }

        // Check if the source has at least one image
        guard CGImageSourceGetCount(imageSource) > 0 else {
            return nil
        }

        // Get image properties to determine orientation
        let options: [CFString: Any] = [
            kCGImageSourceShouldCache: false,
            kCGImageSourceShouldAllowFloat: true
        ]

        guard let cgImage = CGImageSourceCreateImageAtIndex(imageSource, 0, options as CFDictionary) else {
            return nil
        }

        // Extract EXIF orientation if available
        let orientation = extractOrientation(from: imageSource)

        // Apply orientation correction
        let correctedImage = applyOrientationCorrection(to: cgImage, orientation: orientation)

        return processImage(correctedImage)
    }

    // MARK: - Private Methods

    /// Extracts CGImage from UIImage, handling orientation.
    private func extractCGImage(from image: UIImage) -> CGImage? {
        // If the image already has correct orientation (.up), return cgImage directly
        if image.imageOrientation == .up {
            return image.cgImage
        }

        // Otherwise, render the image with correct orientation
        return renderWithCorrectOrientation(image)
    }

    /// Renders a UIImage to a new CGImage with correct orientation applied.
    private func renderWithCorrectOrientation(_ image: UIImage) -> CGImage? {
        let size = image.size

        guard size.width > 0 && size.height > 0 else {
            return nil
        }

        // Create a bitmap context with sRGB color space
        guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB) else {
            // Fallback to device RGB if sRGB is unavailable
            guard let deviceColorSpace = CGColorSpace(name: CGColorSpace.genericRGBLinear) else {
                return nil
            }
            return renderImage(image, size: size, colorSpace: deviceColorSpace)
        }

        return renderImage(image, size: size, colorSpace: colorSpace)
    }

    /// Renders image to a new bitmap context.
    private func renderImage(_ image: UIImage, size: CGSize, colorSpace: CGColorSpace) -> CGImage? {
        let width = Int(size.width)
        let height = Int(size.height)

        guard width > 0 && height > 0 else {
            return nil
        }

        let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue)

        guard let context = CGContext(
            data: nil,
            width: width,
            height: height,
            bitsPerComponent: 8,
            bytesPerRow: 0,
            space: colorSpace,
            bitmapInfo: bitmapInfo.rawValue
        ) else {
            return nil
        }

        // UIKit draws with origin at top-left, CGContext at bottom-left
        context.translateBy(x: 0, y: CGFloat(height))
        context.scaleBy(x: 1.0, y: -1.0)

        UIGraphicsPushContext(context)
        image.draw(in: CGRect(origin: .zero, size: size))
        UIGraphicsPopContext()

        return context.makeImage()
    }

    /// Extracts EXIF orientation from image source.
    private func extractOrientation(from source: CGImageSource) -> CGImagePropertyOrientation {
        guard let properties = CGImageSourceCopyPropertiesAtIndex(source, 0, nil) as? [CFString: Any],
              let orientationValue = properties[kCGImagePropertyOrientation] as? UInt32,
              let orientation = CGImagePropertyOrientation(rawValue: orientationValue) else {
            return .up
        }
        return orientation
    }

    /// Applies orientation correction to a CGImage.
    private func applyOrientationCorrection(
        to image: CGImage,
        orientation: CGImagePropertyOrientation
    ) -> CGImage {
        // If orientation is already correct, return as-is
        guard orientation != .up else {
            return image
        }

        let width = image.width
        let height = image.height

        // Calculate the size of the output image
        let outputSize: (width: Int, height: Int)
        switch orientation {
        case .left, .leftMirrored, .right, .rightMirrored:
            outputSize = (height, width)
        default:
            outputSize = (width, height)
        }

        // Create context with sRGB color space
        guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB) ?? image.colorSpace else {
            return image
        }

        let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue)

        guard let context = CGContext(
            data: nil,
            width: outputSize.width,
            height: outputSize.height,
            bitsPerComponent: 8,
            bytesPerRow: 0,
            space: colorSpace,
            bitmapInfo: bitmapInfo.rawValue
        ) else {
            return image
        }

        // Apply transform based on orientation
        applyTransform(to: context, for: orientation, width: width, height: height)

        // Draw the image
        context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))

        return context.makeImage() ?? image
    }

    /// Applies the appropriate transform to the context based on EXIF orientation.
    private func applyTransform(
        to context: CGContext,
        for orientation: CGImagePropertyOrientation,
        width: Int,
        height: Int
    ) {
        let w = CGFloat(width)
        let h = CGFloat(height)

        switch orientation {
        case .up:
            break
        case .upMirrored:
            context.translateBy(x: w, y: 0)
            context.scaleBy(x: -1, y: 1)
        case .down:
            context.translateBy(x: w, y: h)
            context.rotate(by: .pi)
        case .downMirrored:
            context.translateBy(x: 0, y: h)
            context.scaleBy(x: 1, y: -1)
        case .left:
            context.translateBy(x: h, y: 0)
            context.rotate(by: .pi / 2)
        case .leftMirrored:
            context.translateBy(x: h, y: w)
            context.rotate(by: .pi / 2)
            context.scaleBy(x: 1, y: -1)
        case .right:
            context.translateBy(x: 0, y: w)
            context.rotate(by: -.pi / 2)
        case .rightMirrored:
            context.translateBy(x: 0, y: 0)
            context.rotate(by: -.pi / 2)
            context.scaleBy(x: 1, y: -1)
        }
    }

    /// Processes a CGImage for ML model input.
    private func processImage(_ image: CGImage) -> CGImage? {
        // Validate dimensions
        guard validateDimensions(width: image.width, height: image.height) else {
            return nil
        }

        // Convert to sRGB if required
        if configuration.requiresSRGB {
            return convertToSRGB(image)
        }

        return image
    }

    /// Validates that image dimensions meet minimum requirements.
    private func validateDimensions(width: Int, height: Int) -> Bool {
        return width >= configuration.minimumWidth && height >= configuration.minimumHeight
    }

    /// Converts a CGImage to sRGB color space if needed.
    private func convertToSRGB(_ image: CGImage) -> CGImage? {
        // Check if already in sRGB
        if let colorSpace = image.colorSpace,
           colorSpace.name == CGColorSpace.sRGB {
            return image
        }

        // Create sRGB color space
        guard let srgbColorSpace = CGColorSpace(name: CGColorSpace.sRGB) else {
            return image // Return original if sRGB is unavailable
        }

        let width = image.width
        let height = image.height

        let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue)

        guard let context = CGContext(
            data: nil,
            width: width,
            height: height,
            bitsPerComponent: 8,
            bytesPerRow: 0,
            space: srgbColorSpace,
            bitmapInfo: bitmapInfo.rawValue
        ) else {
            return image // Return original if context creation fails
        }

        context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))

        return context.makeImage() ?? image
    }
}

// MARK: - Throwing Variant

extension ImagePreprocessor {
    /// Prepares a UIImage for ML model input, throwing detailed errors on failure.
    /// - Parameter image: The source UIImage to prepare.
    /// - Returns: A CGImage ready for Vision framework processing.
    /// - Throws: `ImagePreprocessorError` if preparation fails.
    func prepareWithValidation(image: UIImage) throws -> CGImage {
        guard let cgImage = extractCGImage(from: image) else {
            throw ImagePreprocessorError.cgImageCreationFailed
        }

        // Validate dimensions
        guard validateDimensions(width: cgImage.width, height: cgImage.height) else {
            throw ImagePreprocessorError.dimensionsTooSmall(
                width: cgImage.width,
                height: cgImage.height
            )
        }

        // Convert to sRGB if required
        if configuration.requiresSRGB {
            guard let srgbImage = convertToSRGB(cgImage) else {
                throw ImagePreprocessorError.colorSpaceConversionFailed
            }
            return srgbImage
        }

        return cgImage
    }

    /// Prepares image data for ML model input, throwing detailed errors on failure.
    /// - Parameter data: Raw image data (JPEG, PNG, HEIC, etc.).
    /// - Returns: A CGImage ready for Vision framework processing.
    /// - Throws: `ImagePreprocessorError` if preparation fails.
    func prepareWithValidation(data: Data) throws -> CGImage {
        guard !data.isEmpty else {
            throw ImagePreprocessorError.corruptData
        }

        guard let imageSource = CGImageSourceCreateWithData(data as CFData, nil) else {
            throw ImagePreprocessorError.corruptData
        }

        guard CGImageSourceGetCount(imageSource) > 0 else {
            throw ImagePreprocessorError.unsupportedFormat
        }

        // Check if the format is supported
        guard let type = CGImageSourceGetType(imageSource),
              isFormatSupported(type) else {
            throw ImagePreprocessorError.unsupportedFormat
        }

        let options: [CFString: Any] = [
            kCGImageSourceShouldCache: false,
            kCGImageSourceShouldAllowFloat: true
        ]

        guard let cgImage = CGImageSourceCreateImageAtIndex(imageSource, 0, options as CFDictionary) else {
            throw ImagePreprocessorError.cgImageCreationFailed
        }

        // Extract and apply orientation
        let orientation = extractOrientation(from: imageSource)
        let correctedImage = applyOrientationCorrection(to: cgImage, orientation: orientation)

        // Validate dimensions
        guard validateDimensions(width: correctedImage.width, height: correctedImage.height) else {
            throw ImagePreprocessorError.dimensionsTooSmall(
                width: correctedImage.width,
                height: correctedImage.height
            )
        }

        // Convert to sRGB if required
        if configuration.requiresSRGB {
            guard let srgbImage = convertToSRGB(correctedImage) else {
                throw ImagePreprocessorError.colorSpaceConversionFailed
            }
            return srgbImage
        }

        return correctedImage
    }

    /// Checks if the image format is supported.
    private func isFormatSupported(_ typeIdentifier: CFString) -> Bool {
        let supportedTypes: Set<String> = [
            "public.jpeg",
            "public.png",
            "public.heic",
            "public.heif",
            "com.compuserve.gif",
            "public.tiff",
            "com.microsoft.bmp",
            "com.apple.icns"
        ]

        return supportedTypes.contains(typeIdentifier as String)
    }
}