to-live-photo/Sources/LivePhotoCore/AIEnhancer/ImageFormatConverter.swift

//
//  ImageFormatConverter.swift
//  LivePhotoCore
//
//  Utilities for converting between CGImage and CVPixelBuffer formats.
//

import Accelerate
import CoreGraphics
import CoreVideo
import Foundation
import VideoToolbox

/// Utilities for image format conversion
enum ImageFormatConverter {
    /// Convert CGImage to CVPixelBuffer for Core ML input
    /// - Parameters:
    ///   - image: Input CGImage
    ///   - pixelFormat: Output pixel format (default BGRA)
    /// - Returns: CVPixelBuffer ready for Core ML
    static func cgImageToPixelBuffer(
        _ image: CGImage,
        pixelFormat: OSType = kCVPixelFormatType_32BGRA
    ) throws -> CVPixelBuffer {
        let width = image.width
        let height = image.height

        // Create pixel buffer
        var pixelBuffer: CVPixelBuffer?
        let attrs: [CFString: Any] = [
            kCVPixelBufferCGImageCompatibilityKey: true,
            kCVPixelBufferCGBitmapContextCompatibilityKey: true,
            kCVPixelBufferMetalCompatibilityKey: true,
        ]

        let status = CVPixelBufferCreate(
            kCFAllocatorDefault,
            width,
            height,
            pixelFormat,
            attrs as CFDictionary,
            &pixelBuffer
        )

        guard status == kCVReturnSuccess, let buffer = pixelBuffer else {
            throw AIEnhanceError.inputImageInvalid
        }

        // Lock buffer for writing
        CVPixelBufferLockBaseAddress(buffer, [])
        defer { CVPixelBufferUnlockBaseAddress(buffer, []) }

        guard let baseAddress = CVPixelBufferGetBaseAddress(buffer) else {
            throw AIEnhanceError.inputImageInvalid
        }

        let bytesPerRow = CVPixelBufferGetBytesPerRow(buffer)
        let colorSpace = CGColorSpaceCreateDeviceRGB()

        // Create bitmap context to draw into pixel buffer
        guard
            let context = CGContext(
                data: baseAddress,
                width: width,
                height: height,
                bitsPerComponent: 8,
                bytesPerRow: bytesPerRow,
                space: colorSpace,
                bitmapInfo: CGImageAlphaInfo.premultipliedFirst.rawValue
                    | CGBitmapInfo.byteOrder32Little.rawValue
            )
        else {
            throw AIEnhanceError.inputImageInvalid
        }

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

        return buffer
    }

    /// Convert CVPixelBuffer to CGImage
    /// - Parameter pixelBuffer: Input pixel buffer
    /// - Returns: CGImage representation
    static func pixelBufferToCGImage(_ pixelBuffer: CVPixelBuffer) throws -> CGImage {
        var cgImage: CGImage?
        VTCreateCGImageFromCVPixelBuffer(pixelBuffer, options: nil, imageOut: &cgImage)

        guard let image = cgImage else {
            throw AIEnhanceError.inferenceError("Failed to create CGImage from pixel buffer")
        }

        return image
    }

    /// Extract raw RGBA pixel data from CVPixelBuffer
    /// - Parameter pixelBuffer: Input pixel buffer
    /// - Returns: Array of RGBA bytes
    static func pixelBufferToRGBAData(_ pixelBuffer: CVPixelBuffer) throws -> [UInt8] {
        CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
        defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly) }

        let width = CVPixelBufferGetWidth(pixelBuffer)
        let height = CVPixelBufferGetHeight(pixelBuffer)
        let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)

        guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else {
            throw AIEnhanceError.inferenceError("Cannot access pixel buffer data")
        }

        let pixelFormat = CVPixelBufferGetPixelFormatType(pixelBuffer)

        // Handle BGRA format (most common from Core ML)
        if pixelFormat == kCVPixelFormatType_32BGRA {
            return convertBGRAToRGBA(
                baseAddress: baseAddress,
                width: width,
                height: height,
                bytesPerRow: bytesPerRow
            )
        }

        // Handle RGBA format
        if pixelFormat == kCVPixelFormatType_32RGBA {
            var result = [UInt8](repeating: 0, count: width * height * 4)
            for y in 0..<height {
                let srcRow = baseAddress.advanced(by: y * bytesPerRow)
                let dstOffset = y * width * 4
                memcpy(&result[dstOffset], srcRow, width * 4)
            }
            return result
        }

        // Handle ARGB format
        if pixelFormat == kCVPixelFormatType_32ARGB {
            return convertARGBToRGBA(
                baseAddress: baseAddress,
                width: width,
                height: height,
                bytesPerRow: bytesPerRow
            )
        }

        throw AIEnhanceError.inferenceError("Unsupported pixel format: \(pixelFormat)")
    }

    /// Create CVPixelBuffer from raw RGBA data
    /// - Parameters:
    ///   - rgbaData: RGBA pixel data
    ///   - width: Image width
    ///   - height: Image height
    /// - Returns: CVPixelBuffer
    static func rgbaDataToPixelBuffer(
        _ rgbaData: [UInt8],
        width: Int,
        height: Int
    ) throws -> CVPixelBuffer {
        var pixelBuffer: CVPixelBuffer?
        let attrs: [CFString: Any] = [
            kCVPixelBufferCGImageCompatibilityKey: true,
            kCVPixelBufferCGBitmapContextCompatibilityKey: true,
        ]

        let status = CVPixelBufferCreate(
            kCFAllocatorDefault,
            width,
            height,
            kCVPixelFormatType_32BGRA,
            attrs as CFDictionary,
            &pixelBuffer
        )

        guard status == kCVReturnSuccess, let buffer = pixelBuffer else {
            throw AIEnhanceError.inputImageInvalid
        }

        CVPixelBufferLockBaseAddress(buffer, [])
        defer { CVPixelBufferUnlockBaseAddress(buffer, []) }

        guard let baseAddress = CVPixelBufferGetBaseAddress(buffer) else {
            throw AIEnhanceError.inputImageInvalid
        }

        let bytesPerRow = CVPixelBufferGetBytesPerRow(buffer)

        // Convert RGBA to BGRA while copying
        for y in 0..<height {
            let dstRow = baseAddress.advanced(by: y * bytesPerRow).assumingMemoryBound(to: UInt8.self)
            let srcOffset = y * width * 4

            for x in 0..<width {
                let srcIdx = srcOffset + x * 4
                let dstIdx = x * 4

                // RGBA -> BGRA swap
                dstRow[dstIdx + 0] = rgbaData[srcIdx + 2]  // B
                dstRow[dstIdx + 1] = rgbaData[srcIdx + 1]  // G
                dstRow[dstIdx + 2] = rgbaData[srcIdx + 0]  // R
                dstRow[dstIdx + 3] = rgbaData[srcIdx + 3]  // A
            }
        }

        return buffer
    }

    // MARK: - Private Helpers

    private static func convertBGRAToRGBA(
        baseAddress: UnsafeMutableRawPointer,
        width: Int,
        height: Int,
        bytesPerRow: Int
    ) -> [UInt8] {
        var result = [UInt8](repeating: 0, count: width * height * 4)

        for y in 0..<height {
            let srcRow = baseAddress.advanced(by: y * bytesPerRow).assumingMemoryBound(to: UInt8.self)
            let dstOffset = y * width * 4

            for x in 0..<width {
                let srcIdx = x * 4
                let dstIdx = dstOffset + x * 4

                // BGRA -> RGBA swap
                result[dstIdx + 0] = srcRow[srcIdx + 2]  // R
                result[dstIdx + 1] = srcRow[srcIdx + 1]  // G
                result[dstIdx + 2] = srcRow[srcIdx + 0]  // B
                result[dstIdx + 3] = srcRow[srcIdx + 3]  // A
            }
        }

        return result
    }

    private static func convertARGBToRGBA(
        baseAddress: UnsafeMutableRawPointer,
        width: Int,
        height: Int,
        bytesPerRow: Int
    ) -> [UInt8] {
        var result = [UInt8](repeating: 0, count: width * height * 4)

        for y in 0..<height {
            let srcRow = baseAddress.advanced(by: y * bytesPerRow).assumingMemoryBound(to: UInt8.self)
            let dstOffset = y * width * 4

            for x in 0..<width {
                let srcIdx = x * 4
                let dstIdx = dstOffset + x * 4

                // ARGB -> RGBA swap
                result[dstIdx + 0] = srcRow[srcIdx + 1]  // R
                result[dstIdx + 1] = srcRow[srcIdx + 2]  // G
                result[dstIdx + 2] = srcRow[srcIdx + 3]  // B
                result[dstIdx + 3] = srcRow[srcIdx + 0]  // A
            }
        }

        return result
    }
}