//
//  EdgeRefinement.swift
//  CheapRetouch
//
//  Smart brush edge detection using gradient magnitude analysis.
//

import Foundation
import CoreGraphics
import Accelerate
import UIKit

struct EdgeRefinement: Sendable {

    struct GradientImage: Sendable {
        let width: Int
        let height: Int
        let magnitude: [Float]
        let directionX: [Float]
        let directionY: [Float]
    }

    nonisolated static func computeGradient(from image: CGImage) -> GradientImage? {
        let width = image.width
        let height = image.height

        // Convert to grayscale
        guard let grayscale = convertToGrayscaleNonisolated(image) else {
            return nil
        }

        // Compute Sobel gradients
        var gradientX = [Float](repeating: 0, count: width * height)
        var gradientY = [Float](repeating: 0, count: width * height)
        var magnitude = [Float](repeating: 0, count: width * height)

        // Sobel kernels
        let sobelX: [Int16] = [-1, 0, 1, -2, 0, 2, -1, 0, 1]
        let sobelY: [Int16] = [-1, -2, -1, 0, 0, 0, 1, 2, 1]

        // Apply Sobel filters using direct calculation
        for y in 1..<(height - 1) {
            for x in 1..<(width - 1) {
                var gx: Float = 0
                var gy: Float = 0

                for ky in -1...1 {
                    for kx in -1...1 {
                        let pixel = Float(grayscale[(y + ky) * width + (x + kx)])
                        let kernelIndex = (ky + 1) * 3 + (kx + 1)
                        gx += pixel * Float(sobelX[kernelIndex])
                        gy += pixel * Float(sobelY[kernelIndex])
                    }
                }

                let index = y * width + x
                gradientX[index] = gx
                gradientY[index] = gy
                magnitude[index] = sqrt(gx * gx + gy * gy)
            }
        }

        return GradientImage(
            width: width,
            height: height,
            magnitude: magnitude,
            directionX: gradientX,
            directionY: gradientY
        )
    }

    nonisolated static func refineSelectionToEdges(
        selection: [CGPoint],
        gradient: GradientImage,
        searchRadius: Int = 5
    ) -> [CGPoint] {
        var refinedPoints: [CGPoint] = []

        for point in selection {
            let x = Int(point.x)
            let y = Int(point.y)

            // Find the nearby point with highest gradient magnitude
            var maxMagnitude: Float = 0
            var bestX = x
            var bestY = y

            for dy in -searchRadius...searchRadius {
                for dx in -searchRadius...searchRadius {
                    let nx = x + dx
                    let ny = y + dy

                    guard nx >= 0, nx < gradient.width,
                          ny >= 0, ny < gradient.height else { continue }

                    let index = ny * gradient.width + nx
                    let mag = gradient.magnitude[index]

                    if mag > maxMagnitude {
                        maxMagnitude = mag
                        bestX = nx
                        bestY = ny
                    }
                }
            }

            refinedPoints.append(CGPoint(x: bestX, y: bestY))
        }

        return refinedPoints
    }

    nonisolated static func createMaskFromPoints(
        _ points: [CGPoint],
        brushSize: CGFloat,
        imageSize: CGSize
    ) -> CGImage? {
        let width = Int(imageSize.width)
        let height = Int(imageSize.height)

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

        context.setFillColor(gray: 0, alpha: 1)
        context.fill(CGRect(x: 0, y: 0, width: width, height: height))

        context.setFillColor(gray: 1, alpha: 1)

        for point in points {
            let rect = CGRect(
                x: point.x - brushSize / 2,
                y: point.y - brushSize / 2,
                width: brushSize,
                height: brushSize
            )
            context.fillEllipse(in: rect)
        }

        return context.makeImage()
    }

    private nonisolated static func convertToGrayscaleNonisolated(_ image: CGImage) -> [UInt8]? {
        let width = image.width
        let height = image.height
        var pixelData = [UInt8](repeating: 0, count: width * height)

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

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