import SwiftUI
import Metal
import MetalKit
import simd

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(_ view: MTKView, drawableSizeWillChange size: 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()
        }
    }
}