// Open license, I don't care what you do with this.

// idk how many of these are needed
import SwiftUI
import RealityKit
import Metal
import MetalKit
import Spatial

class DrawableWrapper {
    var wrapped: AnyObject? = nil
    var drawable: AnyObject? = nil
    var textureResource: TextureResource? = nil
    
    init(pixelFormat: MTLPixelFormat, width: Int, height: Int, usage: MTLTextureUsage) {
        if #available(visionOS 2.0, *) {
            let desc = LowLevelTexture.Descriptor(textureType: .type2D, pixelFormat: pixelFormat, width: width, height: height, depth: 1, mipmapLevelCount: 1, arrayLength: 1, textureUsage: usage)
            let tex = try? LowLevelTexture(descriptor: desc)
            self.wrapped = tex
        }
        else {
            let desc = TextureResource.DrawableQueue.Descriptor(pixelFormat: pixelFormat, width: width, height: height, usage: [usage], mipmapsMode: .none)
            let queue = try? TextureResource.DrawableQueue(desc)
            queue!.allowsNextDrawableTimeout = true
            self.wrapped = queue
        }
    }
    
    func makeTextureResource() -> TextureResource? {
        if #available(visionOS 2.0, *) {
            if let tex = wrapped as? LowLevelTexture {
                self.textureResource = try! TextureResource(from: tex)
                return self.textureResource
            }
        }
        
        if let queue = wrapped as? TextureResource.DrawableQueue {
            if self.textureResource == nil {
                let data = Data([0x00, 0x00, 0x00, 0xFF])
                self.textureResource = try! TextureResource(
                    dimensions: .dimensions(width: 1, height: 1),
                    format: .raw(pixelFormat: .bgra8Unorm),
                    contents: .init(
                        mipmapLevels: [
                            .mip(data: data, bytesPerRow: 4),
                        ]
                    )
                )
            }
            self.textureResource?.replace(withDrawables: queue)
            
            return self.textureResource
        }
        
        return nil
    }
    
    func nextTexture(commandBuffer: MTLCommandBuffer) -> MTLTexture? {
        if #available(visionOS 2.0, *) {
            if let tex = wrapped as? LowLevelTexture {
                let writeTexture: MTLTexture = tex.replace(using: commandBuffer)
                return writeTexture
            }
        }
        
        if let queue = wrapped as? TextureResource.DrawableQueue {
            self.drawable = try? queue.nextDrawable()
            return self.drawable?.texture
        }
        
        return nil
    }
    
    @MainActor func present(commandBuffer: MTLCommandBuffer) {
        if #available(visionOS 2.0, *) {
            if let tex = wrapped as? LowLevelTexture {
                commandBuffer.commit()
                //commandBuffer.waitUntilCompleted()
                return
            }
        }
        if let drawable = self.drawable as? TextureResource.Drawable {
            commandBuffer.commit()
            commandBuffer.waitUntilCompleted()
            drawable.presentOnSceneUpdate()
        }
    }
}