//
//  MaskData.swift
//  CheapRetouch
//
//  Mask data utilities for compression and decompression.
//

import Foundation
import CoreGraphics
import UIKit
import Accelerate

struct MaskData: Sendable {
    let width: Int
    let height: Int
    let data: Data

    init(width: Int, height: Int, data: Data) {
        self.width = width
        self.height = height
        self.data = data
    }

    init?(from cgImage: CGImage) {
        let width = cgImage.width
        let height = cgImage.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.draw(cgImage, in: CGRect(x: 0, y: 0, width: width, height: height))

        guard let data = context.data else {
            return nil
        }

        let buffer = Data(bytes: data, count: width * height)

        self.width = width
        self.height = height
        self.data = buffer
    }

    func toCGImage() -> CGImage? {
        var mutableData = data
        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,
                space: CGColorSpaceCreateDeviceGray(),
                bitmapInfo: CGImageAlphaInfo.none.rawValue
            ) else {
                return nil
            }

            return context.makeImage()
        }
    }

    func dilated(by pixels: Int) -> MaskData? {
        guard pixels > 0 else { return self }

        let count = width * height
        var sourceArray = [UInt8](data)
        var destinationArray = [UInt8](repeating: 0, count: count)

        let kernelSize = pixels * 2 + 1
        let kernel = [UInt8](repeating: 255, count: kernelSize * kernelSize)

        let error = sourceArray.withUnsafeMutableBufferPointer { sourcePtr -> vImage_Error in
            destinationArray.withUnsafeMutableBufferPointer { destPtr -> vImage_Error in
                var sourceBuffer = vImage_Buffer(
                    data: sourcePtr.baseAddress,
                    height: vImagePixelCount(height),
                    width: vImagePixelCount(width),
                    rowBytes: width
                )

                var destinationBuffer = vImage_Buffer(
                    data: destPtr.baseAddress,
                    height: vImagePixelCount(height),
                    width: vImagePixelCount(width),
                    rowBytes: width
                )

                return vImageDilate_Planar8(
                    &sourceBuffer,
                    &destinationBuffer,
                    0, 0,
                    kernel,
                    vImagePixelCount(kernelSize),
                    vImagePixelCount(kernelSize),
                    vImage_Flags(kvImageNoFlags)
                )
            }
        }

        guard error == kvImageNoError else { return nil }

        return MaskData(width: width, height: height, data: Data(destinationArray))
    }
}
