//
//  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
    private let previewDiffusionIterations: Int = 30
    private let fullDiffusionIterations: Int = 500

    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 {
        DebugLogger.processing("InpaintEngine.inpaint called")
        DebugLogger.log("Image: \(image.width)x\(image.height), Mask: \(mask.width)x\(mask.height)")

        // Check memory requirements
        let imageMemory = image.width * image.height * 4
        let maskMemory = mask.width * mask.height
        let estimatedTotal = imageMemory * 3 + maskMemory * 2 // rough estimate

        DebugLogger.log("Estimated memory: \(estimatedTotal / 1_000_000)MB, max: \(maxMemoryBytes / 1_000_000)MB")

        guard estimatedTotal < maxMemoryBytes else {
            DebugLogger.error("Memory pressure - image too large")
            throw InpaintError.memoryPressure
        }

        // Try LaMa AI inpainting first for best quality
        do {
            DebugLogger.log("Attempting LaMa AI inpainting...")
            let lamaResult = try await lamaInpainter.inpaint(image: image, mask: mask)
            DebugLogger.imageInfo("Inpaint result (LaMa)", image: lamaResult)
            return lamaResult
        } catch {
            DebugLogger.log("LaMa failed: \(error.localizedDescription), falling back to gradient fill")
        }

        // Fallback to gradient-based Metal inpainting
        if isMetalAvailable {
            DebugLogger.log("Using Metal gradient fill for inpainting")
            return try await inpaintWithMetal(image: image, mask: mask, isPreview: false)
        } else {
            DebugLogger.log("Using Accelerate fallback for inpainting")
            return try await inpaintWithAccelerate(image: image, mask: mask)
        }
    }
    
    /// LaMa inpainter for AI-powered inpainting (cached on first use)
    private var _lamaInpainter: LaMaInpainter?
    private var lamaInpainter: LaMaInpainter {
        get async {
            if let existing = _lamaInpainter {
                return existing
            }
            let inpainter = await MainActor.run { LaMaInpainter() }
            _lamaInpainter = inpainter
            return inpainter
        }
    }

    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 else {
            throw InpaintError.metalNotAvailable
        }

        // Create inpainter with appropriate iteration count for preview vs full
        let iterations = isPreview ? previewDiffusionIterations : fullDiffusionIterations
        let featherAmount: Float = isPreview ? 2.0 : 4.0

        do {
            // Metal operations need to run on main actor due to GPU resource management
            let result = try await MainActor.run {
                let inpainter = try PatchMatchInpainter(
                    device: device,
                    patchRadius: 8,
                    diffusionIterations: iterations
                )
                return try inpainter.inpaint(image: image, mask: mask, featherAmount: featherAmount)
            }
            return result
        } catch {
            // If Metal inpainting fails, fall back to Accelerate
            print("Metal inpainting failed: \(error.localizedDescription), falling back to Accelerate")
            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
        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
    }
}
