
import SwiftUI
import Metal
import MetalKit
import UIKit

public protocol CircleMaskedAuraViewCoordinatorDelegate: AnyObject {
    func didUpdateAuraFrame(_ deltaTime: TimeInterval)
}

public struct CircleMaskedAuraView: UIViewRepresentable {
    
    @ObservedObject var coordinator: Coordinator
    let frameRate: Int
    
    public init(_ coordinator: Coordinator,
        morphSpeed: ADFloat = 0.1,
        scale: ADFloat = 1.0,
        frameRate: Int = 60) {
        self.frameRate = frameRate
        self.coordinator = coordinator
        coordinator.setGradientSpeed(morphSpeed)
        coordinator.setScale(scale)
    }
    
    public func makeCoordinator() -> Coordinator {
        return coordinator
    }
    
    public func makeUIView(context: Context) -> some UIView {
        let mtkView = MTKView()
        mtkView.delegate = context.coordinator
        mtkView.preferredFramesPerSecond = frameRate
        mtkView.enableSetNeedsDisplay = true
        mtkView.device = Adamantium.sharedDevice
        
        mtkView.framebufferOnly = false
        mtkView.clearColor = MTLClearColor(red: 0, green: 0, blue: 0, alpha: 0)
        mtkView.drawableSize = mtkView.frame.size
        mtkView.isPaused = false
        return mtkView
    }
    
    public func updateUIView(_ uiView: UIViewType, context: Context) {
        
    }
    
    public class Coordinator: NSObject, ObservableObject, MTKViewDelegate {
        
        private let id = UUID()
        var metalDevice: MTLDevice?
        var metalCommandQueue: MTLCommandQueue?
        public weak var delegate: CircleMaskedAuraViewCoordinatorDelegate?
        
        public struct PropertyTarget {
            var mainColorA: ADVector3 = .zero
            var mainColorB: ADVector3 = .zero
            var mainColorC: ADVector3 = .zero
            
            var darkColorA: ADVector3 = .zero
            var darkColorB: ADVector3 = .zero
            var darkColorC: ADVector3 = .zero
            
            var maskPosX: ADFloat = 0.5
            var maskPosY: ADFloat = 0.5
            public var maskEdgeFade: ADFloat = 0.3
            var maskRadius: ADFloat = 0.25
            var maskMinValueScalar: ADFloat = 0.3
            
            public init(
                mainColorA: ADVector3,
                mainColorB: ADVector3,
                mainColorC: ADVector3,
                darkColorA: ADVector3,
                darkColorB: ADVector3,
                darkColorC: ADVector3,
                maskPosX: ADFloat,
                maskPosY: ADFloat,
                maskEdgeFade: ADFloat,
                maskMinValueScalar: ADFloat,
                maskRadius: ADFloat) {
                    
                self.mainColorA = mainColorA
                self.mainColorB = mainColorB
                self.mainColorC = mainColorC
                self.darkColorA = darkColorA
                self.darkColorB = darkColorB
                self.darkColorC = darkColorC
                self.maskPosX = maskPosX
                self.maskPosY = maskPosY
                self.maskEdgeFade = maskEdgeFade
                self.maskMinValueScalar = maskMinValueScalar
                self.maskRadius = maskRadius
            }
            
            public func mix(_ mixTarget: PropertyTarget, ratio: Float) -> PropertyTarget {
                let newMainColorA           = ADVector3.lerpVector(mainColorA, target: mixTarget.mainColorA, ratio: ratio)
                let newMainColorB           = ADVector3.lerpVector(mainColorB, target: mixTarget.mainColorB, ratio: ratio)
                let newMainColorC           = ADVector3.lerpVector(mainColorC, target: mixTarget.mainColorC, ratio: ratio)
                
                let newDarkColorA           = ADVector3.lerpVector(darkColorA, target: mixTarget.darkColorA, ratio: ratio)
                let newDarkColorB           = ADVector3.lerpVector(darkColorB, target: mixTarget.darkColorB, ratio: ratio)
                let newDarkColorC           = ADVector3.lerpVector(darkColorC, target: mixTarget.darkColorC, ratio: ratio)
                
                let newMaskPosX             = simd_mix(maskPosX, mixTarget.maskPosX, ratio)
                let newMaskPosY             = simd_mix(maskPosY, mixTarget.maskPosY, ratio)
                let newMaskEdgeFade         = simd_mix(maskEdgeFade, mixTarget.maskEdgeFade, ratio)
                let newMaskMinValueScalar   = simd_mix(maskMinValueScalar, mixTarget.maskMinValueScalar, ratio)
                let newMaskRadius           = simd_mix(maskRadius, mixTarget.maskRadius, ratio)
                
                return .init(
                    mainColorA: newMainColorA,
                    mainColorB: newMainColorB,
                    mainColorC: newMainColorC,
                    darkColorA: newDarkColorA,
                    darkColorB: newDarkColorB,
                    darkColorC: newDarkColorC,
                    maskPosX: newMaskPosX,
                    maskPosY: newMaskPosY,
                    maskEdgeFade: newMaskEdgeFade,
                    maskMinValueScalar: newMaskMinValueScalar,
                    maskRadius: newMaskRadius
                )
            }
        }
        
        private var propertyTarget: PropertyTarget?
        private var propertyTargetOverride: PropertyTarget?
        private var targetAnimationSpeedScalar: ADFloat = 1.0
        
        private var lastMainColorA: SIMD3<Float> = .zero
        private var lastMainColorB: SIMD3<Float> = .zero
        private var lastMainColorC: SIMD3<Float> = .zero
        
        private var lastDarkColorA: SIMD3<Float> = .zero
        private var lastDarkColorB: SIMD3<Float> = .zero
        private var lastDarkColorC: SIMD3<Float> = .zero
        
        private var xPosMask: ADFloat = 0.5
        private var yPosMask: ADFloat = 0.5
        private var maskEdgeFade: ADFloat = 0.3
        private var lastCircleRadius: ADFloat = 0.25
        private var lastMinValueScalar: ADFloat = 0.3
        
        private var gradientSpeed: ADFloat = 0.1
        private var gradientScale: ADFloat = 1.0
        private var gradientSeed: ADFloat = .zero
        
        private let startingTime: Date = .now
        private var lastTime: Date?
        
        public override init() {
            self.metalDevice = Adamantium.sharedDevice
            self.metalCommandQueue = Adamantium.sharedCommandQueue
            super.init()
        }
        
        public init(animationTimeScalar: ADFloat) {
            self.metalDevice = Adamantium.sharedDevice
            self.metalCommandQueue = Adamantium.sharedCommandQueue
            super.init()
            
            targetAnimationSpeedScalar = animationTimeScalar
        }
        
        public func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) { }
        
        func setGradientSpeed(_ speedScalar: ADFloat) {
            gradientSpeed = speedScalar
        }
        
        func setScale(_ scale: ADFloat) {
            gradientScale = scale
        }
        
        public func setGradientSeed(_ seed: ADFloat = .zero) {
            gradientSeed = seed
        }
        
        func setLastCircleRadius(_ circleRadius: ADFloat) {
            lastCircleRadius = circleRadius
        }
        
        func setMaskPosX(_ position: ADFloat) {
            xPosMask = position
        }
        
        func setMaskPosY(_ position: ADFloat) {
            yPosMask = position
        }
        
        func setMaskEdgeFade(_ fade: ADFloat) {
            maskEdgeFade = fade
        }
        
        public func setAnimationTimeScalar(_ scalar: TimeInterval) {
            targetAnimationSpeedScalar = ADFloat(scalar)
        }
        
        public func applyPropertiesWithoutTransitionAnimation(_ target: PropertyTarget) {
            propertyTargetOverride = target
        }
        
        public func setPropertyTarget(_ target: PropertyTarget) {
            propertyTarget = target
        }
        
        public func draw(in view: MTKView) {
            var size: ADVector2 {
                return ADVector2(
                    x: Float(view.drawableSize.width),
                    y: Float(view.drawableSize.height)
                )
            }
            
            let deltaTime: TimeInterval
            if let lastTime = lastTime {
                deltaTime = Date.now.timeIntervalSince(lastTime)
            } else {
                deltaTime = .zero
            }
            
            lastTime = .now
            animateToTarget(deltaTime)
            
            guard
                let commandBuffer = Adamantium.sharedCommandQueue?.makeCommandBuffer(),
                let currentDrawable = view.currentDrawable,
                let renderDescriptor = view.currentRenderPassDescriptor,
                let renderEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderDescriptor)
            else { return }

            delegate?.didUpdateAuraFrame(deltaTime)
            renderAura(size, renderEncoder: renderEncoder)
            
            renderEncoder.endEncoding()
            commandBuffer.present(currentDrawable)
            commandBuffer.commit()
        }
    }
}

private extension CircleMaskedAuraView.Coordinator {
    
    func lerpColorToTarget(_ color: ADVector3, target: ADVector3, deltaTime: Float) -> ADVector3 {
        return ADVector3.lerpVector(color, target: target, ratio: deltaTime)
    }
    
    func animateToTarget(_ deltaTime: TimeInterval) {
        if let overrideTarget = propertyTargetOverride {
            lastMainColorA      = overrideTarget.mainColorA
            lastMainColorB      = overrideTarget.mainColorB
            lastMainColorC      = overrideTarget.mainColorC
            
            lastDarkColorA      = overrideTarget.darkColorA
            lastDarkColorB      = overrideTarget.darkColorB
            lastDarkColorC      = overrideTarget.darkColorC
            
            lastCircleRadius    = overrideTarget.maskRadius
            xPosMask            = overrideTarget.maskPosX
            yPosMask            = overrideTarget.maskPosY
            maskEdgeFade        = overrideTarget.maskEdgeFade
            lastMinValueScalar  = overrideTarget.maskMinValueScalar
            
            propertyTargetOverride = nil
            
        } else {
            let timeScalar = targetAnimationSpeedScalar
            guard let target = propertyTarget else { return }
            let floatDeltaTime = Float(deltaTime) * timeScalar
            
            lastMainColorA      = lerpColorToTarget(lastMainColorA, target: target.mainColorA, deltaTime: floatDeltaTime)
            lastMainColorB      = lerpColorToTarget(lastMainColorB, target: target.mainColorB, deltaTime: floatDeltaTime)
            lastMainColorC      = lerpColorToTarget(lastMainColorC, target: target.mainColorC, deltaTime: floatDeltaTime)
            
            lastDarkColorA      = lerpColorToTarget(lastDarkColorA, target: target.darkColorA, deltaTime: floatDeltaTime)
            lastDarkColorB      = lerpColorToTarget(lastDarkColorB, target: target.darkColorB, deltaTime: floatDeltaTime)
            lastDarkColorC      = lerpColorToTarget(lastDarkColorC, target: target.darkColorC, deltaTime: floatDeltaTime)
            
            lastCircleRadius    = simd_mix(lastCircleRadius,    target.maskRadius, floatDeltaTime)
            xPosMask            = simd_mix(xPosMask,            target.maskPosX, floatDeltaTime)
            yPosMask            = simd_mix(yPosMask,            target.maskPosY, floatDeltaTime)
            maskEdgeFade        = simd_mix(maskEdgeFade,        target.maskEdgeFade, floatDeltaTime)
            lastMinValueScalar  = simd_mix(lastMinValueScalar,  target.maskMinValueScalar, floatDeltaTime)
        }
    }
    
    func renderAura(_ size: ADVector2, renderEncoder: MTLRenderCommandEncoder) {
        let timeElapsed = ADFloat(Date.now.timeIntervalSince(startingTime))
        let fractalElapsedTime = timeElapsed * gradientSpeed
        let grainElapsedTime = timeElapsed * 0.01
        
        let noiseRenderable = MaskedCircleAuraRenderable()
        noiseRenderable.config = .init(
            displaySize: size,
            fractalNoiseX: .zero,
            fractalNoiseY: .zero,
            fractalNoiseZ: fractalElapsedTime + gradientSeed,
            fractalUniformScale: 0.5 * gradientScale,
            grainNoiseZ: grainElapsedTime,
            grainNoiseScale: 5.0 * gradientScale,
            
            maskRadius: lastCircleRadius,
            maskCenterX: xPosMask,
            maskCenterY: yPosMask,
            maskEdgeFade: maskEdgeFade,
            maskMinValueScalar: lastMinValueScalar,
            
            backgroundColor: SunoColorPalette.backgroundDotCom.floatThree,
            
            // Main color
            colorA: lastMainColorA,
            colorB: lastMainColorB,
            colorC: lastMainColorC,
            
            // Dark color
            darkColorA: lastDarkColorA,
            darkColorB: lastDarkColorB,
            darkColorC: lastDarkColorC,
            
            // Dark mixing
            darkMixPointStart: -0.2,
            darkMixPointBackground: -0.1,
            darkMixPointBackgroundA: 0.0,
            darkMixPointAB: 0.2,
            darkMixPointBC: 0.4,
            darkMixPointEnd: 2.0,
            
            // Color mixing
            mixPointStart: 0.0,
            mixPointBackground: 0.1,
            mixPointBackgroundA: 0.2,
            mixPointAB: 0.3,
            mixPointBC: 0.4,
            mixPointEnd: 2.0
        )
        
        noiseRenderable.render(encoder: renderEncoder)
    }
}
