//
//  PatchMatch.swift
//  CheapRetouch
//
//  Exemplar-based inpainting algorithm implementation using Metal.
//

import Foundation
import Metal
import MetalKit
import CoreGraphics
import simd

final class PatchMatchInpainter {

    private let device: MTLDevice
    private let commandQueue: MTLCommandQueue
    private let library: MTLLibrary

    // Compute pipelines
    private let dilateMaskPipeline: MTLComputePipelineState
    private let gaussianBlurPipeline: MTLComputePipelineState
    private let diffuseInpaintPipeline: MTLComputePipelineState
    private let edgeAwareBlendPipeline: MTLComputePipelineState
    private let gradientFillPipeline: MTLComputePipelineState

    private let patchRadius: Int
    private let diffusionIterations: Int

    init(device: MTLDevice, patchRadius: Int = 4, diffusionIterations: Int = 100) throws {
        self.device = device
        self.patchRadius = patchRadius
        self.diffusionIterations = diffusionIterations

        guard let queue = device.makeCommandQueue() else {
            throw PatchMatchError.commandQueueCreationFailed
        }
        self.commandQueue = queue

        guard let library = device.makeDefaultLibrary() else {
            throw PatchMatchError.libraryNotFound
        }
        self.library = library

        // Create compute pipelines
        guard let dilateFunc = library.makeFunction(name: "dilateMask") else {
            throw PatchMatchError.functionNotFound("dilateMask")
        }
        self.dilateMaskPipeline = try device.makeComputePipelineState(function: dilateFunc)

        guard let blurFunc = library.makeFunction(name: "gaussianBlur") else {
            throw PatchMatchError.functionNotFound("gaussianBlur")
        }
        self.gaussianBlurPipeline = try device.makeComputePipelineState(function: blurFunc)

        guard let diffuseFunc = library.makeFunction(name: "diffuseInpaint") else {
            throw PatchMatchError.functionNotFound("diffuseInpaint")
        }
        self.diffuseInpaintPipeline = try device.makeComputePipelineState(function: diffuseFunc)

        guard let blendFunc = library.makeFunction(name: "edgeAwareBlend") else {
            throw PatchMatchError.functionNotFound("edgeAwareBlend")
        }
        self.edgeAwareBlendPipeline = try device.makeComputePipelineState(function: blendFunc)
        
        guard let gradientFunc = library.makeFunction(name: "gradientFill") else {
            throw PatchMatchError.functionNotFound("gradientFill")
        }
        self.gradientFillPipeline = try device.makeComputePipelineState(function: gradientFunc)
    }

    func inpaint(image: CGImage, mask: CGImage, featherAmount: Float = 4.0) throws -> CGImage {
        let width = image.width
        let height = image.height

        DebugLogger.processing("PatchMatchInpainter.inpaint started")
        DebugLogger.log("Creating textures: \(width)x\(height), mask: \(mask.width)x\(mask.height)")
        DebugLogger.log("Patch radius: \(patchRadius), diffusion iterations: \(diffusionIterations), feather: \(featherAmount)")

        // Create textures
        let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .rgba8Unorm,
            width: width,
            height: height,
            mipmapped: false
        )
        textureDescriptor.usage = [.shaderRead, .shaderWrite]

        guard let sourceTexture = device.makeTexture(descriptor: textureDescriptor),
              let resultTexture = device.makeTexture(descriptor: textureDescriptor),
              let tempTexture = device.makeTexture(descriptor: textureDescriptor) else {
            throw PatchMatchError.textureCreationFailed
        }

        // Mask texture (single channel)
        let maskDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .r8Unorm,
            width: width,
            height: height,
            mipmapped: false
        )
        maskDescriptor.usage = [.shaderRead, .shaderWrite]

        guard let maskTexture = device.makeTexture(descriptor: maskDescriptor),
              let dilatedMaskTexture = device.makeTexture(descriptor: maskDescriptor),
              let featheredMaskTexture = device.makeTexture(descriptor: maskDescriptor) else {
            throw PatchMatchError.textureCreationFailed
        }

        // Load image data into texture
        DebugLogger.log("Loading source image into texture...")
        try loadCGImage(image, into: sourceTexture)
        DebugLogger.log("Loading mask into texture (will scale from \(mask.width)x\(mask.height) to \(width)x\(height))...")
        try loadCGImageGrayscale(mask, into: maskTexture)
        DebugLogger.log("Textures loaded successfully")

        // Copy source to result for initial state
        try copyTexture(from: sourceTexture, to: resultTexture)

        guard let commandBuffer = commandQueue.makeCommandBuffer() else {
            throw PatchMatchError.commandBufferCreationFailed
        }

        // Step 1: Dilate mask significantly to exclude object edge pixels
        // Use 20px dilation to ensure we sample from clean background, not object edges
        let dilationRadius = max(patchRadius, 20)
        DebugLogger.log("Dilating mask by \(dilationRadius) pixels...")
        encodeDilateMask(commandBuffer: commandBuffer, input: maskTexture, output: dilatedMaskTexture, radius: dilationRadius)

        // Step 2: Feather mask
        encodeGaussianBlur(commandBuffer: commandBuffer, input: dilatedMaskTexture, output: featheredMaskTexture, radius: Int(featherAmount))

        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()

        // Step 3: Gradient-based fill (much faster than diffusion for uniform backgrounds)
        // This samples colors from 8 directions and blends based on distance
        DebugLogger.log("Starting gradient fill...")
        guard let gradientBuffer = commandQueue.makeCommandBuffer() else {
            throw PatchMatchError.commandBufferCreationFailed
        }
        
        encodeGradientFill(
            commandBuffer: gradientBuffer,
            source: sourceTexture,
            mask: dilatedMaskTexture,
            output: resultTexture,
            maxDistance: Float(max(width, height))
        )
        
        gradientBuffer.commit()
        gradientBuffer.waitUntilCompleted()
        DebugLogger.log("Gradient fill complete")

        // Step 4: Apply diffusion iterations to smooth the result and eliminate seams
        let smoothingIterations = min(diffusionIterations, 100)
        DebugLogger.log("Applying \(smoothingIterations) smoothing iterations...")
        for _ in 0..<smoothingIterations {
            guard let iterBuffer = commandQueue.makeCommandBuffer() else {
                throw PatchMatchError.commandBufferCreationFailed
            }

            encodeDiffuseInpaint(commandBuffer: iterBuffer, input: resultTexture, mask: dilatedMaskTexture, output: tempTexture)

            iterBuffer.commit()
            iterBuffer.waitUntilCompleted()

            try copyTexture(from: tempTexture, to: resultTexture)
        }
        DebugLogger.log("Smoothing complete")

        // Step 5: Edge-aware blending
        guard let finalBuffer = commandQueue.makeCommandBuffer() else {
            throw PatchMatchError.commandBufferCreationFailed
        }

        encodeEdgeAwareBlend(
            commandBuffer: finalBuffer,
            inpainted: resultTexture,
            original: sourceTexture,
            mask: featheredMaskTexture,
            output: tempTexture,
            blendRadius: featherAmount
        )

        finalBuffer.commit()
        finalBuffer.waitUntilCompleted()

        // Convert result to CGImage
        return try extractCGImage(from: tempTexture)
    }

    // MARK: - Shader Encoding

    private func encodeDilateMask(commandBuffer: MTLCommandBuffer, input: MTLTexture, output: MTLTexture, radius: Int) {
        guard let encoder = commandBuffer.makeComputeCommandEncoder() else { return }

        encoder.setComputePipelineState(dilateMaskPipeline)
        encoder.setTexture(input, index: 0)
        encoder.setTexture(output, index: 1)

        var radiusValue = Int32(radius)
        encoder.setBytes(&radiusValue, length: MemoryLayout<Int32>.size, index: 0)

        let threadgroupSize = MTLSize(width: 16, height: 16, depth: 1)
        let threadgroups = MTLSize(
            width: (input.width + 15) / 16,
            height: (input.height + 15) / 16,
            depth: 1
        )

        encoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadgroupSize)
        encoder.endEncoding()
    }

    private func encodeGaussianBlur(commandBuffer: MTLCommandBuffer, input: MTLTexture, output: MTLTexture, radius: Int) {
        // Generate Gaussian weights
        let sigma = Float(radius) / 3.0
        var weights: [Float] = []
        for i in 0...radius {
            let weight = exp(-Float(i * i) / (2 * sigma * sigma))
            weights.append(weight)
        }

        guard let weightsBuffer = device.makeBuffer(bytes: weights, length: weights.count * MemoryLayout<Float>.size, options: []) else { return }

        // Create temp texture for two-pass blur
        let tempDescriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: input.pixelFormat,
            width: input.width,
            height: input.height,
            mipmapped: false
        )
        tempDescriptor.usage = [.shaderRead, .shaderWrite]

        guard let tempTexture = device.makeTexture(descriptor: tempDescriptor) else { return }

        // Horizontal pass
        if let encoder = commandBuffer.makeComputeCommandEncoder() {
            encoder.setComputePipelineState(gaussianBlurPipeline)
            encoder.setTexture(input, index: 0)
            encoder.setTexture(tempTexture, index: 1)
            encoder.setBuffer(weightsBuffer, offset: 0, index: 0)
            var radiusValue = Int32(radius)
            encoder.setBytes(&radiusValue, length: MemoryLayout<Int32>.size, index: 1)
            var horizontal: Bool = true
            encoder.setBytes(&horizontal, length: MemoryLayout<Bool>.size, index: 2)

            let threadgroupSize = MTLSize(width: 16, height: 16, depth: 1)
            let threadgroups = MTLSize(
                width: (input.width + 15) / 16,
                height: (input.height + 15) / 16,
                depth: 1
            )

            encoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadgroupSize)
            encoder.endEncoding()
        }

        // Vertical pass
        if let encoder = commandBuffer.makeComputeCommandEncoder() {
            encoder.setComputePipelineState(gaussianBlurPipeline)
            encoder.setTexture(tempTexture, index: 0)
            encoder.setTexture(output, index: 1)
            encoder.setBuffer(weightsBuffer, offset: 0, index: 0)
            var radiusValue = Int32(radius)
            encoder.setBytes(&radiusValue, length: MemoryLayout<Int32>.size, index: 1)
            var horizontal: Bool = false
            encoder.setBytes(&horizontal, length: MemoryLayout<Bool>.size, index: 2)

            let threadgroupSize = MTLSize(width: 16, height: 16, depth: 1)
            let threadgroups = MTLSize(
                width: (input.width + 15) / 16,
                height: (input.height + 15) / 16,
                depth: 1
            )

            encoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadgroupSize)
            encoder.endEncoding()
        }
    }

    private func encodeDiffuseInpaint(commandBuffer: MTLCommandBuffer, input: MTLTexture, mask: MTLTexture, output: MTLTexture) {
        guard let encoder = commandBuffer.makeComputeCommandEncoder() else { return }

        encoder.setComputePipelineState(diffuseInpaintPipeline)
        encoder.setTexture(input, index: 0)
        encoder.setTexture(mask, index: 1)
        encoder.setTexture(output, index: 2)

        let threadgroupSize = MTLSize(width: 16, height: 16, depth: 1)
        let threadgroups = MTLSize(
            width: (input.width + 15) / 16,
            height: (input.height + 15) / 16,
            depth: 1
        )

        encoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadgroupSize)
        encoder.endEncoding()
    }

    private func encodeEdgeAwareBlend(
        commandBuffer: MTLCommandBuffer,
        inpainted: MTLTexture,
        original: MTLTexture,
        mask: MTLTexture,
        output: MTLTexture,
        blendRadius: Float
    ) {
        guard let encoder = commandBuffer.makeComputeCommandEncoder() else { return }

        encoder.setComputePipelineState(edgeAwareBlendPipeline)
        encoder.setTexture(inpainted, index: 0)
        encoder.setTexture(original, index: 1)
        encoder.setTexture(mask, index: 2)
        encoder.setTexture(output, index: 3)

        var radius = blendRadius
        encoder.setBytes(&radius, length: MemoryLayout<Float>.size, index: 0)

        let threadgroupSize = MTLSize(width: 16, height: 16, depth: 1)
        let threadgroups = MTLSize(
            width: (inpainted.width + 15) / 16,
            height: (inpainted.height + 15) / 16,
            depth: 1
        )

        encoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadgroupSize)
        encoder.endEncoding()
    }
    
    private func encodeGradientFill(
        commandBuffer: MTLCommandBuffer,
        source: MTLTexture,
        mask: MTLTexture,
        output: MTLTexture,
        maxDistance: Float
    ) {
        guard let encoder = commandBuffer.makeComputeCommandEncoder() else { return }

        encoder.setComputePipelineState(gradientFillPipeline)
        encoder.setTexture(source, index: 0)
        encoder.setTexture(mask, index: 1)
        // We don't have a distance texture, so we'll pass mask twice (distanceTexture is unused in our simplified version)
        encoder.setTexture(mask, index: 2)
        encoder.setTexture(output, index: 3)

        var maxDist = maxDistance
        encoder.setBytes(&maxDist, length: MemoryLayout<Float>.size, index: 0)

        let threadgroupSize = MTLSize(width: 16, height: 16, depth: 1)
        let threadgroups = MTLSize(
            width: (source.width + 15) / 16,
            height: (source.height + 15) / 16,
            depth: 1
        )

        encoder.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadgroupSize)
        encoder.endEncoding()
    }

    // MARK: - Texture Utilities

    private func loadCGImage(_ cgImage: CGImage, into texture: MTLTexture) throws {
        // Use texture dimensions for consistency
        let width = texture.width
        let height = texture.height

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

        // Draw the image scaled to fit the texture dimensions
        context.interpolationQuality = .high
        context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))

        guard let data = context.data else {
            throw PatchMatchError.dataExtractionFailed
        }

        texture.replace(
            region: MTLRegionMake2D(0, 0, width, height),
            mipmapLevel: 0,
            withBytes: data,
            bytesPerRow: width * 4
        )
    }

    private func loadCGImageGrayscale(_ cgImage: CGImage, into texture: MTLTexture) throws {
        // Use texture dimensions, scaling the mask if needed
        let width = texture.width
        let height = texture.height

        guard let context = CGContext(
            data: nil,
            width: width,
            height: height,
            bitsPerComponent: 8,
            bytesPerRow: width,
            space: CGColorSpaceCreateDeviceGray(),
            bitmapInfo: CGImageAlphaInfo.none.rawValue
        ) else {
            throw PatchMatchError.contextCreationFailed
        }

        // Draw the mask scaled to fit the texture dimensions
        context.interpolationQuality = .high
        context.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))

        guard let data = context.data else {
            throw PatchMatchError.dataExtractionFailed
        }
        
        // Debug: count non-zero pixels to verify mask has content
        let buffer = data.bindMemory(to: UInt8.self, capacity: width * height)
        var whitePixelCount = 0
        var totalValue: Int = 0
        for i in 0..<(width * height) {
            let value = buffer[i]
            totalValue += Int(value)
            if value > 127 {
                whitePixelCount += 1
            }
        }
        DebugLogger.log("Mask loaded: \(width)x\(height), whitePixels=\(whitePixelCount) (\(String(format: "%.2f", Double(whitePixelCount) / Double(width * height) * 100))%), avgValue=\(totalValue / (width * height))")

        texture.replace(
            region: MTLRegionMake2D(0, 0, width, height),
            mipmapLevel: 0,
            withBytes: data,
            bytesPerRow: width
        )
    }

    private func copyTexture(from source: MTLTexture, to destination: MTLTexture) throws {
        guard let commandBuffer = commandQueue.makeCommandBuffer(),
              let blitEncoder = commandBuffer.makeBlitCommandEncoder() else {
            throw PatchMatchError.commandBufferCreationFailed
        }

        blitEncoder.copy(
            from: source,
            sourceSlice: 0,
            sourceLevel: 0,
            sourceOrigin: MTLOrigin(x: 0, y: 0, z: 0),
            sourceSize: MTLSize(width: source.width, height: source.height, depth: 1),
            to: destination,
            destinationSlice: 0,
            destinationLevel: 0,
            destinationOrigin: MTLOrigin(x: 0, y: 0, z: 0)
        )

        blitEncoder.endEncoding()
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()
    }

    private func extractCGImage(from texture: MTLTexture) throws -> CGImage {
        let width = texture.width
        let height = texture.height
        let bytesPerRow = width * 4

        var pixelData = [UInt8](repeating: 0, count: bytesPerRow * height)

        texture.getBytes(
            &pixelData,
            bytesPerRow: bytesPerRow,
            from: MTLRegionMake2D(0, 0, width, height),
            mipmapLevel: 0
        )

        guard let context = CGContext(
            data: &pixelData,
            width: width,
            height: height,
            bitsPerComponent: 8,
            bytesPerRow: bytesPerRow,
            space: CGColorSpaceCreateDeviceRGB(),
            bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue
        ) else {
            throw PatchMatchError.contextCreationFailed
        }

        guard let cgImage = context.makeImage() else {
            throw PatchMatchError.imageCreationFailed
        }

        return cgImage
    }
}

// MARK: - Errors

enum PatchMatchError: LocalizedError {
    case commandQueueCreationFailed
    case libraryNotFound
    case functionNotFound(String)
    case textureCreationFailed
    case commandBufferCreationFailed
    case contextCreationFailed
    case dataExtractionFailed
    case imageCreationFailed

    var errorDescription: String? {
        switch self {
        case .commandQueueCreationFailed:
            return "Failed to create Metal command queue"
        case .libraryNotFound:
            return "Metal shader library not found"
        case .functionNotFound(let name):
            return "Metal function '\(name)' not found"
        case .textureCreationFailed:
            return "Failed to create Metal texture"
        case .commandBufferCreationFailed:
            return "Failed to create command buffer"
        case .contextCreationFailed:
            return "Failed to create graphics context"
        case .dataExtractionFailed:
            return "Failed to extract image data"
        case .imageCreationFailed:
            return "Failed to create output image"
        }
    }
}
