import Metal
import MetalKit
import simd
import SwiftUI

public struct WarpedGridView: UIViewRepresentable, Identifiable {
    public let id: String

    // Exposed uniforms:
    public var timeScale: ADFloat
    public var colorA: ADVector3
    public var colorB: ADVector3
    public var vortexRadius: ADFloat
    public var vortexStrength: ADFloat

    // Frame control:
    public var frameRate: Int = 60
    public var isPaused: Bool = false

    @ObservedObject private var coordinator: Coordinator

    public init(
        id: String,
        timeScale: ADFloat = 1.0,
        colorA: ADVector3 = ADVector3(1.0, 0.3, 0.5),
        colorB: ADVector3 = ADVector3(1.0, 0.45, 0.4),
        vortexRadius: ADFloat = 4.0,
        vortexStrength: ADFloat = -0.4,
        frameRate: Int = 60,
        isPaused: Bool = false
    ) {
        self.id = id
        self.timeScale = timeScale
        self.colorA = colorA
        self.colorB = colorB
        self.vortexRadius = vortexRadius
        self.vortexStrength = vortexStrength
        self.frameRate = frameRate
        self.isPaused = isPaused

        coordinator = Coordinator(
            timeScale: timeScale,
            colorA: colorA,
            colorB: colorB,
            vortexRadius: vortexRadius,
            vortexStrength: vortexStrength
        )
    }

    public func makeCoordinator() -> Coordinator { coordinator }

    public func makeUIView(context: Context) -> MTKView {
        let mtk = MTKView(frame: .zero, device: MTLCreateSystemDefaultDevice())
        mtk.delegate = context.coordinator
        mtk.preferredFramesPerSecond = frameRate
        mtk.framebufferOnly = false
        mtk.clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 0)
        mtk.enableSetNeedsDisplay = true
        mtk.isPaused = isPaused
        return mtk
    }

    public func updateUIView(_ uiView: MTKView, context _: Context) {
        uiView.isPaused = isPaused
        coordinator.update(
            timeScale: timeScale,
            colorA: colorA,
            colorB: colorB,
            vortexRadius: vortexRadius,
            vortexStrength: vortexStrength
        )
    }

    public class Coordinator: NSObject, ObservableObject, MTKViewDelegate {
        private let device: MTLDevice?
        private let queue: MTLCommandQueue?
        private let renderable: WarpedGridRenderable

        private var uniforms = WarpedGridRenderable.Config.defaultConfiguration
        private let start = CFAbsoluteTimeGetCurrent()
        private var lastTime: CFAbsoluteTime?
        private var timeScale: ADFloat
        private var colorA: ADVector3
        private var colorB: ADVector3
        private var vortexRadius: ADFloat
        private var vortexStrength: ADFloat

        public init(
            timeScale: ADFloat,
            colorA: ADVector3,
            colorB: ADVector3,
            vortexRadius: ADFloat,
            vortexStrength: ADFloat
        ) {
            self.device = Adamantium.sharedDevice
            self.queue = Adamantium.sharedCommandQueue
            self.renderable = WarpedGridRenderable()
            self.timeScale = timeScale
            self.colorA = colorA
            self.colorB = colorB
            self.vortexRadius = vortexRadius
            self.vortexStrength = vortexStrength
            super.init()
        }

        func update(
            timeScale: ADFloat,
            colorA: ADVector3,
            colorB: ADVector3,
            vortexRadius: ADFloat,
            vortexStrength: ADFloat
        ) {
            self.timeScale = timeScale
            self.colorA = colorA
            self.colorB = colorB
            self.vortexRadius = vortexRadius
            self.vortexStrength = vortexStrength
        }

        public func mtkView(_: MTKView, drawableSizeWillChange _: CGSize) {
            // nothing special
        }

        public func draw(in view: MTKView) {
            guard
                let buf = queue?.makeCommandBuffer(),
                let desc = view.currentRenderPassDescriptor,
                let drw = view.currentDrawable
            else { return }

            let now = CFAbsoluteTimeGetCurrent()
            let delta = lastTime.map { now - $0 } ?? 0
            lastTime = now
            let elapsed = Float(now - start)

            // update uniforms
            uniforms.resolution = ADVector2(
                Float(view.drawableSize.width),
                Float(view.drawableSize.height)
            )
            uniforms.time = 10 + elapsed * timeScale
            uniforms.colorA = colorA
            uniforms.colorB = colorB
            uniforms.vortexRadius = vortexRadius
            uniforms.vortexStrength = vortexStrength

            renderable.config = uniforms

            let enc = buf.makeRenderCommandEncoder(descriptor: desc)!
            renderable.render(encoder: enc)
            enc.endEncoding()

            buf.present(drw)
            buf.commit()
        }
    }
}