CheapRetouch/CheapRetouch/Services/InpaintEngine/InpaintEngine.swift

//
//  InpaintEngine.swift
//  CheapRetouch
//
//  Exemplar-based inpainting engine using Metal.
//

import Foundation
import Metal
import MetalKit
import CoreGraphics
import UIKit
import Accelerate

actor InpaintEngine {

    enum InpaintError: Error, LocalizedError {
        case metalNotAvailable
        case deviceCreationFailed
        case textureCreationFailed
        case processingFailed
        case memoryPressure
        case invalidInput

        var errorDescription: String? {
            switch self {
            case .metalNotAvailable:
                return "Metal is not available on this device"
            case .deviceCreationFailed:
                return "Failed to create Metal device"
            case .textureCreationFailed:
                return "Failed to create texture for processing"
            case .processingFailed:
                return "Inpainting processing failed"
            case .memoryPressure:
                return "Image too large to process. Try cropping first."
            case .invalidInput:
                return "Invalid image or mask input"
            }
        }
    }

    private let device: MTLDevice?
    private let commandQueue: MTLCommandQueue?
    private let patchRadius: Int
    private let maxPreviewSize: Int = 2048
    private let maxMemoryBytes: Int = 1_500_000_000 // 1.5GB

    init(patchRadius: Int = 9) {
        self.patchRadius = patchRadius
        self.device = MTLCreateSystemDefaultDevice()
        self.commandQueue = device?.makeCommandQueue()
    }

    var isMetalAvailable: Bool {
        device != nil && commandQueue != nil
    }

    func inpaint(image: CGImage, mask: CGImage) async throws -> CGImage {
        // Check memory requirements
        let imageMemory = image.width * image.height * 4
        let maskMemory = mask.width * mask.height
        let estimatedTotal = imageMemory * 3 + maskMemory * 2 // rough estimate

        guard estimatedTotal < maxMemoryBytes else {
            throw InpaintError.memoryPressure
        }

        if isMetalAvailable {
            return try await inpaintWithMetal(image: image, mask: mask, isPreview: false)
        } else {
            return try await inpaintWithAccelerate(image: image, mask: mask)
        }
    }

    func inpaintPreview(image: CGImage, mask: CGImage) async throws -> CGImage {
        // Scale down for preview if needed
        let scaledImage: CGImage
        let scaledMask: CGImage

        if image.width > maxPreviewSize || image.height > maxPreviewSize {
            let scale = CGFloat(maxPreviewSize) / CGFloat(max(image.width, image.height))
            scaledImage = try scaleImage(image, scale: scale)
            scaledMask = try scaleImage(mask, scale: scale)
        } else {
            scaledImage = image
            scaledMask = mask
        }

        if isMetalAvailable {
            return try await inpaintWithMetal(image: scaledImage, mask: scaledMask, isPreview: true)
        } else {
            return try await inpaintWithAccelerate(image: scaledImage, mask: scaledMask)
        }
    }

    // MARK: - Metal Implementation

    private func inpaintWithMetal(image: CGImage, mask: CGImage, isPreview: Bool) async throws -> CGImage {
        guard let device = device, let commandQueue = commandQueue else {
            throw InpaintError.metalNotAvailable
        }

        // For now, fall back to Accelerate while Metal shaders are being developed
        return try await inpaintWithAccelerate(image: image, mask: mask)
    }

    // MARK: - Accelerate Fallback

    private func inpaintWithAccelerate(image: CGImage, mask: CGImage) async throws -> CGImage {
        let width = image.width
        let height = image.height

        guard width > 0, height > 0 else {
            throw InpaintError.invalidInput
        }

        // Convert image to pixel buffer
        guard var imageBuffer = createPixelBuffer(from: image) else {
            throw InpaintError.textureCreationFailed
        }

        // Convert mask to grayscale buffer
        guard let maskBuffer = createGrayscaleBuffer(from: mask, targetWidth: width, targetHeight: height) else {
            throw InpaintError.textureCreationFailed
        }

        // Simple inpainting: for each masked pixel, average nearby unmasked pixels
        let patchSize = patchRadius * 2 + 1

        for y in 0..<height {
            for x in 0..<width {
                let maskIndex = y * width + x
                guard maskBuffer[maskIndex] > 127 else { continue }

                // This pixel needs to be filled
                var totalR: Int = 0
                var totalG: Int = 0
                var totalB: Int = 0
                var count: Int = 0

                // Search in expanding rings for source pixels
                for radius in 1...max(patchRadius * 4, 50) {
                    for dy in -radius...radius {
                        for dx in -radius...radius {
                            // Only check perimeter of current radius
                            guard abs(dx) == radius || abs(dy) == radius else { continue }

                            let sx = x + dx
                            let sy = y + dy

                            guard sx >= 0, sx < width, sy >= 0, sy < height else { continue }

                            let srcMaskIndex = sy * width + sx
                            guard maskBuffer[srcMaskIndex] <= 127 else { continue }

                            let srcPixelIndex = (sy * width + sx) * 4
                            totalR += Int(imageBuffer[srcPixelIndex])
                            totalG += Int(imageBuffer[srcPixelIndex + 1])
                            totalB += Int(imageBuffer[srcPixelIndex + 2])
                            count += 1
                        }
                    }

                    if count >= 8 { break }
                }

                if count > 0 {
                    let pixelIndex = (y * width + x) * 4
                    imageBuffer[pixelIndex] = UInt8(totalR / count)
                    imageBuffer[pixelIndex + 1] = UInt8(totalG / count)
                    imageBuffer[pixelIndex + 2] = UInt8(totalB / count)
                    imageBuffer[pixelIndex + 3] = 255
                }
            }
        }

        // Convert back to CGImage
        guard let result = createCGImage(from: imageBuffer, width: width, height: height) else {
            throw InpaintError.processingFailed
        }

        return result
    }

    // MARK: - Helper Methods

    private func createPixelBuffer(from image: CGImage) -> [UInt8]? {
        let width = image.width
        let height = image.height
        var pixelData = [UInt8](repeating: 0, count: width * height * 4)

        guard let context = CGContext(
            data: &pixelData,
            width: width,
            height: height,
            bitsPerComponent: 8,
            bytesPerRow: width * 4,
            space: CGColorSpaceCreateDeviceRGB(),
            bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
        ) else {
            return nil
        }

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

    private func createGrayscaleBuffer(from image: CGImage, targetWidth: Int, targetHeight: Int) -> [UInt8]? {
        var pixelData = [UInt8](repeating: 0, count: targetWidth * targetHeight)

        guard let context = CGContext(
            data: &pixelData,
            width: targetWidth,
            height: targetHeight,
            bitsPerComponent: 8,
            bytesPerRow: targetWidth,
            space: CGColorSpaceCreateDeviceGray(),
            bitmapInfo: CGImageAlphaInfo.none.rawValue
        ) else {
            return nil
        }

        context.draw(image, in: CGRect(x: 0, y: 0, width: targetWidth, height: targetHeight))
        return pixelData
    }

    private func createCGImage(from pixelData: [UInt8], width: Int, height: Int) -> CGImage? {
        var mutableData = pixelData

        return mutableData.withUnsafeMutableBytes { rawBuffer -> CGImage? in
            guard let baseAddress = rawBuffer.baseAddress else { return nil }

            guard let context = CGContext(
                data: baseAddress,
                width: width,
                height: height,
                bitsPerComponent: 8,
                bytesPerRow: width * 4,
                space: CGColorSpaceCreateDeviceRGB(),
                bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
            ) else {
                return nil
            }

            return context.makeImage()
        }
    }

    private func scaleImage(_ image: CGImage, scale: CGFloat) throws -> CGImage {
        let newWidth = Int(CGFloat(image.width) * scale)
        let newHeight = Int(CGFloat(image.height) * scale)

        guard let context = CGContext(
            data: nil,
            width: newWidth,
            height: newHeight,
            bitsPerComponent: 8,
            bytesPerRow: newWidth * 4,
            space: CGColorSpaceCreateDeviceRGB(),
            bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
        ) else {
            throw InpaintError.textureCreationFailed
        }

        context.interpolationQuality = .high
        context.draw(image, in: CGRect(x: 0, y: 0, width: newWidth, height: newHeight))

        guard let result = context.makeImage() else {
            throw InpaintError.textureCreationFailed
        }

        return result
    }
}