// ViewportRenderer.swift
// Draws all objects from the SceneGraph using Metal.
// Each ViewportPane owns one instance with its own command queue + camera.

import Metal
import MetalKit
import Observation
import simd

@Observable
final class ViewportRenderer {

    // MARK: - Public state

    var viewportSize: SIMD2<Float> = SIMD2<Float>(1, 1)
    let camera: OrbitalCamera

    // Weak ref to the shared scene graph (set externally)
    var sceneGraph: SceneGraph?

    // MARK: - Metal objects (owned per-pane)

    private(set) var device: MTLDevice?
    private var commandQueue: MTLCommandQueue?

    // Pipelines (shared via MetalSharedState)
    private var meshPipeline: MTLRenderPipelineState?
    private var wirePipeline: MTLRenderPipelineState?
    private var gridPipeline: MTLRenderPipelineState?
    private var axisPipeline: MTLRenderPipelineState?
    private var selectionPipeline: MTLRenderPipelineState?

    private var depthState: MTLDepthStencilState?
    private var depthStateNoWrite: MTLDepthStencilState?

    // Axis gizmo buffer
    private var axisVertexBuffer: MTLBuffer?

    init(camera: OrbitalCamera = OrbitalCamera()) {
        self.camera = camera
    }

    // MARK: - Setup

    func configure(device: MTLDevice, library: MTLLibrary) {
        self.device = device
        self.commandQueue = device.makeCommandQueue()!

        // ── Pipelines ──
        meshPipeline = Self.makePipeline(device: device, library: library,
            vertex: "mesh_vertex", fragment: "mesh_fragment",
            label: "Mesh", blending: false)
        wirePipeline = Self.makePipeline(device: device, library: library,
            vertex: "wire_vertex", fragment: "wire_fragment",
            label: "Wire", blending: true)
        gridPipeline = Self.makePipeline(device: device, library: library,
            vertex: "grid_vertex", fragment: "grid_fragment",
            label: "Grid", blending: true)
        axisPipeline = Self.makePipeline(device: device, library: library,
            vertex: "axis_vertex", fragment: "axis_fragment",
            label: "Axis", blending: true)
        selectionPipeline = Self.makePipeline(device: device, library: library,
            vertex: "wire_vertex", fragment: "wire_fragment",
            label: "Selection", blending: true)

        // ── Depth states ──
        let dd = MTLDepthStencilDescriptor()
        dd.depthCompareFunction = .less
        dd.isDepthWriteEnabled = true
        depthState = device.makeDepthStencilState(descriptor: dd)

        let ddNoWrite = MTLDepthStencilDescriptor()
        ddNoWrite.depthCompareFunction = .lessEqual
        ddNoWrite.isDepthWriteEnabled = false
        depthStateNoWrite = device.makeDepthStencilState(descriptor: ddNoWrite)

        // ── Axis gizmo ──
        let axisVerts = Self.buildAxisVertices()
        axisVertexBuffer = device.makeBuffer(
            bytes: axisVerts,
            length: MemoryLayout<MeshVertex>.stride * axisVerts.count,
            options: .storageModeShared
        )
    }

    // MARK: - Per-Frame Draw

    func draw(drawable: CAMetalDrawable, passDescriptor: MTLRenderPassDescriptor) {
        guard let commandQueue, let device,
              let meshPipeline, let wirePipeline,
              let gridPipeline, let axisPipeline,
              let depthState, let depthStateNoWrite
        else { return }

        guard let buffer = commandQueue.makeCommandBuffer(),
              let encoder = buffer.makeRenderCommandEncoder(descriptor: passDescriptor)
        else { return }

        let aspect = viewportSize.x / max(viewportSize.y, 1)
        let projection = MathUtils.perspective(
            fovYRadians: Float.pi / 4.0, aspect: aspect,
            near: 0.1, far: 500.0
        )
        let view = camera.viewMatrix

        var sceneConstants = SceneConstants(
            cameraPosition: camera.eye,
            lightDirection: normalize(SIMD3<Float>(-0.5, -1.0, -0.7)),
            ambientIntensity: 0.25,
            diffuseIntensity: 0.75,
            time: 0, _pad: 0
        )

        // ── 1. Grid ──
        drawGrid(encoder: encoder, pipeline: gridPipeline,
                 depthState: depthStateNoWrite,
                 view: view, projection: projection,
                 sceneConstants: &sceneConstants)

        // ── 2. Axis gizmo ──
        if let axisBuf = axisVertexBuffer {
            var identity = Uniforms(
                modelMatrix: matrix_identity_float4x4,
                viewMatrix: view,
                projectionMatrix: projection,
                normalMatrix: matrix_identity_float4x4
            )
            encoder.setRenderPipelineState(axisPipeline)
            encoder.setDepthStencilState(depthState)
            encoder.setVertexBuffer(axisBuf, offset: 0, index: 0)
            encoder.setVertexBytes(&identity, length: MemoryLayout<Uniforms>.size, index: 1)
            encoder.drawPrimitives(type: .line, vertexStart: 0, vertexCount: 6)
        }

        // ── 3. Scene objects ──
        if let scene = sceneGraph {
            for obj in scene.allObjects where obj.isVisible {
                guard let mesh = obj.meshHandle else { continue }

                let model = obj.worldMatrix
                var uniforms = Uniforms(
                    modelMatrix: model,
                    viewMatrix: view,
                    projectionMatrix: projection,
                    normalMatrix: MathUtils.normalMatrix(from: model)
                )

                // Solid geometry
                if mesh.vertexCount > 0 {
                    encoder.setRenderPipelineState(meshPipeline)
                    encoder.setDepthStencilState(depthState)
                    encoder.setVertexBuffer(mesh.vertexBuffer, offset: 0, index: 0)
                    encoder.setVertexBytes(&uniforms, length: MemoryLayout<Uniforms>.size, index: 1)
                    encoder.setFragmentBytes(&sceneConstants, length: MemoryLayout<SceneConstants>.size, index: 0)
                    encoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: mesh.vertexCount)
                }

                // Wireframe overlay
                if let wireBuf = mesh.wireVertexBuffer, mesh.wireVertexCount > 0 {
                    encoder.setRenderPipelineState(wirePipeline)
                    encoder.setDepthStencilState(depthStateNoWrite)
                    encoder.setVertexBuffer(wireBuf, offset: 0, index: 0)
                    encoder.setVertexBytes(&uniforms, length: MemoryLayout<Uniforms>.size, index: 1)
                    encoder.drawPrimitives(type: .line, vertexStart: 0, vertexCount: mesh.wireVertexCount)
                }

                // Selection outline (orange)
                if obj.isSelected, let wireBuf = mesh.wireVertexBuffer, mesh.wireVertexCount > 0 {
                    // Build selection wire verts from existing wire with orange color
                    drawSelectionOutline(
                        encoder: encoder,
                        pipeline: wirePipeline,
                        depthState: depthStateNoWrite,
                        wireBuffer: wireBuf,
                        wireCount: mesh.wireVertexCount,
                        uniforms: &uniforms,
                        device: device
                    )
                }
            }
        }

        encoder.endEncoding()
        buffer.present(drawable)
        buffer.commit()
    }

    // MARK: - Grid

    private func drawGrid(
        encoder: MTLRenderCommandEncoder,
        pipeline: MTLRenderPipelineState,
        depthState: MTLDepthStencilState,
        view: simd_float4x4,
        projection: simd_float4x4,
        sceneConstants: inout SceneConstants
    ) {
        var identity = Uniforms(
            modelMatrix: matrix_identity_float4x4,
            viewMatrix: view,
            projectionMatrix: projection,
            normalMatrix: matrix_identity_float4x4
        )
        encoder.setRenderPipelineState(pipeline)
        encoder.setDepthStencilState(depthState)
        encoder.setVertexBytes(&identity, length: MemoryLayout<Uniforms>.size, index: 0)

        var gridConstants = GridConstants(
            color: SIMD4<Float>(0.5, 0.5, 0.5, 1),
            gridSpacing: 1.0, gridLineWidth: 1.0,
            gridFadeDistance: 40.0, _pad: 0
        )
        encoder.setFragmentBytes(&gridConstants, length: MemoryLayout<GridConstants>.size, index: 0)
        encoder.setFragmentBytes(&sceneConstants, length: MemoryLayout<SceneConstants>.size, index: 1)
        encoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 6)
    }

    // MARK: - Selection outline

    private func drawSelectionOutline(
        encoder: MTLRenderCommandEncoder,
        pipeline: MTLRenderPipelineState,
        depthState: MTLDepthStencilState,
        wireBuffer: MTLBuffer,
        wireCount: Int,
        uniforms: inout Uniforms,
        device: MTLDevice
    ) {
        // We re-use the wire vertex positions but override color in the shader
        // by drawing the same wire buffer with a slightly expanded model matrix
        // Draw selection wireframe (uses wire buffer colors as-is for now)
        encoder.setRenderPipelineState(pipeline)
        encoder.setDepthStencilState(depthState)
        encoder.setVertexBuffer(wireBuffer, offset: 0, index: 0)
        encoder.setVertexBytes(&uniforms, length: MemoryLayout<Uniforms>.size, index: 1)
        encoder.drawPrimitives(type: .line, vertexStart: 0, vertexCount: wireCount)
    }

    // MARK: - Static helpers

    private static func buildAxisVertices() -> [MeshVertex] {
        let len: Float = 1.5
        let red   = SIMD4<Float>(0.78, 0.14, 0.22, 1.0)
        let green = SIMD4<Float>(0.22, 0.68, 0.28, 1.0)
        let blue  = SIMD4<Float>(0.22, 0.45, 0.78, 1.0)
        return [
            .init(position: .zero, normal: .zero, color: red),
            .init(position: .init(len, 0, 0), normal: .zero, color: red),
            .init(position: .zero, normal: .zero, color: green),
            .init(position: .init(0, len, 0), normal: .zero, color: green),
            .init(position: .zero, normal: .zero, color: blue),
            .init(position: .init(0, 0, len), normal: .zero, color: blue),
        ]
    }

    static func makePipeline(
        device: MTLDevice,
        library: MTLLibrary,
        vertex: String,
        fragment: String,
        label: String,
        blending: Bool
    ) -> MTLRenderPipelineState {
        let desc = MTLRenderPipelineDescriptor()
        desc.label = label
        desc.vertexFunction = library.makeFunction(name: vertex)
        desc.fragmentFunction = library.makeFunction(name: fragment)
        desc.colorAttachments[0].pixelFormat = .bgra8Unorm_srgb
        desc.depthAttachmentPixelFormat = .depth32Float

        if blending {
            desc.colorAttachments[0].isBlendingEnabled = true
            desc.colorAttachments[0].sourceRGBBlendFactor = .sourceAlpha
            desc.colorAttachments[0].destinationRGBBlendFactor = .oneMinusSourceAlpha
            desc.colorAttachments[0].sourceAlphaBlendFactor = .one
            desc.colorAttachments[0].destinationAlphaBlendFactor = .oneMinusSourceAlpha
        }

        return try! device.makeRenderPipelineState(descriptor: desc)
    }
}