import CoreVideo
import Metal
import MetalKit
import UIKit
import simd

struct DitherUniforms {
    var texScale: SIMD2<Float>
    var viewSize: SIMD2<Float>
    var spread: Float
    var contrast: Float
    var brightness: Float
    var paletteCount: UInt32
    var matrixSize: UInt32
    var pixelSize: UInt32
}

final class DitherRenderer: NSObject, ObservableObject, MTKViewDelegate {

    let device: MTLDevice

    @Published var palette: Palette = Palette.gameboy {
        didSet { paletteDirty = true }
    }
    @Published var spread: Float = 1.0
    @Published var contrast: Float = 1.15
    @Published var brightness: Float = 0.0
    @Published var matrixSize: UInt32 = 8
    @Published var pixelSize: UInt32 = 2

    var onCapture: ((UIImage?) -> Void)?

    private let commandQueue: MTLCommandQueue
    private let pipeline: MTLRenderPipelineState
    private var textureCache: CVMetalTextureCache?

    private var paletteBuffer: MTLBuffer?
    private var paletteCount: UInt32 = 0
    private var paletteDirty = true

    private let frameLock = NSLock()
    private var pendingPixelBuffer: CVPixelBuffer?

    private var captureRequested = false

    init?() {
        guard let device = MTLCreateSystemDefaultDevice(),
              let commandQueue = device.makeCommandQueue(),
              let source = DitherRenderer.loadShaderSource(),
              let library = try? device.makeLibrary(source: source, options: nil),
              let vertexFunction = library.makeFunction(name: "ditherVertex"),
              let fragmentFunction = library.makeFunction(name: "ditherFragment") else {
            return nil
        }

        let descriptor = MTLRenderPipelineDescriptor()
        descriptor.vertexFunction = vertexFunction
        descriptor.fragmentFunction = fragmentFunction
        descriptor.colorAttachments[0].pixelFormat = .bgra8Unorm

        guard let pipeline = try? device.makeRenderPipelineState(descriptor: descriptor) else {
            return nil
        }

        self.device = device
        self.commandQueue = commandQueue
        self.pipeline = pipeline
        super.init()

        var cache: CVMetalTextureCache?
        let result = CVMetalTextureCacheCreate(kCFAllocatorDefault, nil, device, nil, &cache)
        guard result == kCVReturnSuccess, let cache else { return nil }
        self.textureCache = cache
    }

    private static func loadShaderSource() -> String? {
        let bundles = [Bundle.main, Bundle(for: DitherRenderer.self)]
        for bundle in bundles {
            if let url = bundle.url(forResource: "DitherShaders", withExtension: "txt"),
               let text = try? String(contentsOf: url, encoding: .utf8) {
                return text
            }
        }
        return nil
    }

    func setPixelBuffer(_ pixelBuffer: CVPixelBuffer) {
        frameLock.lock()
        pendingPixelBuffer = pixelBuffer
        frameLock.unlock()
    }

    func requestCapture() {
        captureRequested = true
    }

    func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {
    }

    func draw(in view: MTKView) {
        let drawableSize = view.drawableSize
        guard drawableSize.width >= 1, drawableSize.height >= 1 else { return }
        guard let cache = textureCache else { return }

        frameLock.lock()
        let buffer = pendingPixelBuffer
        frameLock.unlock()
        guard let pixelBuffer = buffer else { return }

        CVMetalTextureCacheFlush(cache, 0)

        let sourceWidth = CVPixelBufferGetWidth(pixelBuffer)
        let sourceHeight = CVPixelBufferGetHeight(pixelBuffer)
        guard sourceWidth > 0, sourceHeight > 0 else { return }

        var cvTexture: CVMetalTexture?
        let result = CVMetalTextureCacheCreateTextureFromImage(
            kCFAllocatorDefault,
            cache,
            pixelBuffer,
            nil,
            .bgra8Unorm,
            sourceWidth,
            sourceHeight,
            0,
            &cvTexture
        )
        guard result == kCVReturnSuccess,
              let cvTexture,
              let sourceTexture = CVMetalTextureGetTexture(cvTexture) else { return }

        updatePaletteBufferIfNeeded()
        guard let paletteBuffer else { return }

        var uniforms = makeUniforms(
            sourceWidth: sourceWidth,
            sourceHeight: sourceHeight,
            viewWidth: Float(drawableSize.width),
            viewHeight: Float(drawableSize.height)
        )

        if let descriptor = view.currentRenderPassDescriptor,
           let drawable = view.currentDrawable,
           let commandBuffer = commandQueue.makeCommandBuffer(),
           let encoder = commandBuffer.makeRenderCommandEncoder(descriptor: descriptor) {
            encoder.setRenderPipelineState(pipeline)
            encoder.setFragmentTexture(sourceTexture, index: 0)
            encoder.setFragmentBytes(&uniforms, length: MemoryLayout<DitherUniforms>.stride, index: 0)
            encoder.setFragmentBuffer(paletteBuffer, offset: 0, index: 1)
            encoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
            encoder.endEncoding()
            commandBuffer.present(drawable)
            let retained = cvTexture
            commandBuffer.addCompletedHandler { _ in
                _ = retained
            }
            commandBuffer.commit()
        }

        if captureRequested {
            captureRequested = false
            renderCapture(
                source: sourceTexture,
                retaining: cvTexture,
                uniforms: uniforms,
                width: Int(drawableSize.width),
                height: Int(drawableSize.height)
            )
        }
    }

    private func makeUniforms(sourceWidth: Int,
                              sourceHeight: Int,
                              viewWidth: Float,
                              viewHeight: Float) -> DitherUniforms {
        let sourceAspect = Float(sourceWidth) / Float(sourceHeight)
        let viewAspect = viewWidth / viewHeight
        var scale = SIMD2<Float>(1.0, 1.0)
        if sourceAspect > viewAspect {
            scale.x = viewAspect / sourceAspect
        } else if sourceAspect < viewAspect {
            scale.y = sourceAspect / viewAspect
        }

        return DitherUniforms(
            texScale: scale,
            viewSize: SIMD2<Float>(viewWidth, viewHeight),
            spread: spread,
            contrast: contrast,
            brightness: brightness,
            paletteCount: paletteCount,
            matrixSize: matrixSize,
            pixelSize: max(pixelSize, 1)
        )
    }

    private func updatePaletteBufferIfNeeded() {
        guard paletteDirty || paletteBuffer == nil else { return }

        var colors = palette.colors
        if colors.isEmpty {
            colors = [SIMD4<Float>(0, 0, 0, 1), SIMD4<Float>(1, 1, 1, 1)]
        }
        if colors.count > Palette.maximumColors {
            colors = Array(colors.prefix(Palette.maximumColors))
        }

        let length = MemoryLayout<SIMD4<Float>>.stride * colors.count
        guard let buffer = device.makeBuffer(bytes: colors, length: length, options: .storageModeShared) else {
            return
        }
        paletteBuffer = buffer
        paletteCount = UInt32(colors.count)
        paletteDirty = false
    }

    private func renderCapture(source: MTLTexture,
                               retaining cvTexture: CVMetalTexture,
                               uniforms: DitherUniforms,
                               width: Int,
                               height: Int) {
        guard width > 0, height > 0, let paletteBuffer else {
            deliverCapture(nil)
            return
        }

        let descriptor = MTLTextureDescriptor.texture2DDescriptor(
            pixelFormat: .bgra8Unorm,
            width: width,
            height: height,
            mipmapped: false
        )
        descriptor.usage = [.renderTarget, .shaderRead]
        descriptor.storageMode = .shared

        guard let target = device.makeTexture(descriptor: descriptor) else {
            deliverCapture(nil)
            return
        }

        let pass = MTLRenderPassDescriptor()
        pass.colorAttachments[0].texture = target
        pass.colorAttachments[0].loadAction = .clear
        pass.colorAttachments[0].storeAction = .store
        pass.colorAttachments[0].clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 1)

        var localUniforms = uniforms

        guard let commandBuffer = commandQueue.makeCommandBuffer(),
              let encoder = commandBuffer.makeRenderCommandEncoder(descriptor: pass) else {
            deliverCapture(nil)
            return
        }

        encoder.setRenderPipelineState(pipeline)
        encoder.setFragmentTexture(source, index: 0)
        encoder.setFragmentBytes(&localUniforms, length: MemoryLayout<DitherUniforms>.stride, index: 0)
        encoder.setFragmentBuffer(paletteBuffer, offset: 0, index: 1)
        encoder.drawPrimitives(type: .triangleStrip, vertexStart: 0, vertexCount: 4)
        encoder.endEncoding()

        commandBuffer.addCompletedHandler { [weak self] finished in
            _ = cvTexture
            guard let self else { return }
            if finished.error != nil {
                self.deliverCapture(nil)
                return
            }
            let image = DitherRenderer.makeImage(from: target)
            self.deliverCapture(image)
        }
        commandBuffer.commit()
    }

    private func deliverCapture(_ image: UIImage?) {
        DispatchQueue.main.async { [weak self] in
            self?.onCapture?(image)
        }
    }

    private static func makeImage(from texture: MTLTexture) -> UIImage? {
        let width = texture.width
        let height = texture.height
        guard width > 0, height > 0 else { return nil }

        let bytesPerRow = width * 4
        var raw = [UInt8](repeating: 0, count: bytesPerRow * height)
        raw.withUnsafeMutableBytes { pointer in
            guard let base = pointer.baseAddress else { return }
            texture.getBytes(
                base,
                bytesPerRow: bytesPerRow,
                from: MTLRegionMake2D(0, 0, width, height),
                mipmapLevel: 0
            )
        }

        guard let provider = CGDataProvider(data: Data(raw) as CFData) else { return nil }

        let bitmapInfo = CGBitmapInfo(
            rawValue: CGImageAlphaInfo.noneSkipFirst.rawValue | CGBitmapInfo.byteOrder32Little.rawValue
        )

        guard let cgImage = CGImage(
            width: width,
            height: height,
            bitsPerComponent: 8,
            bitsPerPixel: 32,
            bytesPerRow: bytesPerRow,
            space: CGColorSpaceCreateDeviceRGB(),
            bitmapInfo: bitmapInfo,
            provider: provider,
            decode: nil,
            shouldInterpolate: false,
            intent: .defaultIntent
        ) else { return nil }

        return UIImage(cgImage: cgImage)
    }
}