import Metal

public final class VisionTower {
    public let config: VisionConfig
    public let engine: MarkBaseEngine
    public let weights: VisionWeights

    private var qBuffer: MTLBuffer
    private var kBuffer: MTLBuffer
    private var vBuffer: MTLBuffer
    private var attnOutBuffer: MTLBuffer
    private var mlpBuffer: MTLBuffer
    private var tempBuffer: MTLBuffer
    private var normBuffer: MTLBuffer
    private var residualBuffer: MTLBuffer

    public init(config: VisionConfig, engine: MarkBaseEngine, weights: VisionWeights) throws {
        self.config = config
        self.engine = engine
        self.weights = weights

        let device = engine.device
        let maxPatches = 4096
        let hiddenSize = config.hiddenSize
        let intermediateSize = config.intermediateSize

        qBuffer = device.makeBuffer(length: hiddenSize * maxPatches * 4)!
        kBuffer = device.makeBuffer(length: hiddenSize * maxPatches * 4)!
        vBuffer = device.makeBuffer(length: hiddenSize * maxPatches * 4)!
        attnOutBuffer = device.makeBuffer(length: hiddenSize * maxPatches * 4)!
        mlpBuffer = device.makeBuffer(length: intermediateSize * maxPatches * 4)!
        tempBuffer = device.makeBuffer(length: max(hiddenSize, intermediateSize) * maxPatches * 4)!
        normBuffer = device.makeBuffer(length: hiddenSize * maxPatches * 4)!
        residualBuffer = device.makeBuffer(length: hiddenSize * maxPatches * 4)!
    }

    public func forward(patchEmbeddings: MTLBuffer, numPatches: Int, outputBuffer: MTLBuffer) throws {
        var current = patchEmbeddings

        let cmdBuf = engine.commandQueue.makeCommandBuffer()!

        // Input projection: [numPatches, 768] -> [numPatches, 768]
        current = try applyQuantizedMatmul(input: current, weights: weights.inputProj,
                                            seqLen: numPatches, output: tempBuffer, cmdBuf: cmdBuf)

        // Add position embedding
        current = try addPositionEmbedding(input: current, numPatches: numPatches, cmdBuf: cmdBuf)

        // Vision layers (16 layers)
        for layerWeights in weights.layers {
            current = try applyLayer(input: current, weights: layerWeights, numPatches: numPatches, cmdBuf: cmdBuf)
        }

        // Embedding projection: [numPatches, 768] -> [numPatches, 2560]
        try applyEmbeddingProjection(input: current, numPatches: numPatches, output: outputBuffer, cmdBuf: cmdBuf)

        cmdBuf.commit()
        cmdBuf.waitUntilCompleted()
    }

    // ── Quantized matmul (sequence-aware) ─────────────

    private func applyQuantizedMatmul(input: MTLBuffer, weights: QuantizedWeights,
                                       seqLen: Int, output: MTLBuffer,
                                       cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let pso = try engine.pipeline(named: "quantized_matmul")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)

        enc.setBuffer(input, offset: 0, index: 0)
        enc.setBuffer(weights.weight, offset: 0, index: 1)
        enc.setBuffer(weights.scales, offset: 0, index: 2)
        enc.setBuffer(weights.biases, offset: 0, index: 3)
        enc.setBuffer(output, offset: 0, index: 4)

        var inD = UInt32(weights.inDim)
        enc.setBytes(&inD, length: MemoryLayout<UInt32>.size, index: 5)
        var outD = UInt32(weights.outDim)
        enc.setBytes(&outD, length: MemoryLayout<UInt32>.size, index: 6)

        let grid = MTLSize(width: weights.outDim * seqLen, height: 1, depth: 1)
        let tg = engine.threadgroupSize1D(pso, count: max(weights.outDim, seqLen))
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()

        return output
    }

    // ── Position embedding ────────────────────────────

    private func addPositionEmbedding(input: MTLBuffer, numPatches: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let output = normBuffer

        let pso = try engine.pipeline(named: "vision_add_pos_embed")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)

        enc.setBuffer(input, offset: 0, index: 0)
        enc.setBuffer(weights.positionEmbedding, offset: 0, index: 1)
        enc.setBuffer(output, offset: 0, index: 2)

        var hiddenSize = UInt32(config.hiddenSize)
        enc.setBytes(&hiddenSize, length: 4, index: 3)
        var numPatches_ = UInt32(numPatches)
        enc.setBytes(&numPatches_, length: 4, index: 4)

        let grid = MTLSize(width: config.hiddenSize, height: numPatches, depth: 1)
        let tg = engine.threadgroupSize2D(pso, grid: (config.hiddenSize, numPatches))
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()

        return output
    }

    // ── Layer ─────────────────────────────────────────

    private func applyLayer(input: MTLBuffer, weights: VisionLayerWeights, numPatches: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        var current = input

        // 1. Input layernorm
        current = try applyRMSNorm(input: current, weight: weights.inputLayernorm, seqLen: numPatches, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf)

        // 2. Self-attention with Q/K norm
        let attnOut = try applyVisionAttention(input: current, weights: weights, numPatches: numPatches, cmdBuf: cmdBuf)

        // 3. Residual + post_attention_layernorm
        current = try applyResidualAdd(input: input, add: attnOut, seqLen: numPatches, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf)
        current = try applyRMSNorm(input: current, weight: weights.postAttentionLayernorm, seqLen: numPatches, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf)

        // 4. Pre-feedforward layernorm
        current = try applyRMSNorm(input: current, weight: weights.preFeedforwardLayernorm, seqLen: numPatches, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf)

        // 5. MLP (SwiGLU)
        let mlpOut = try applyVisionMLP(input: current, weights: weights, numPatches: numPatches, cmdBuf: cmdBuf)

        // 6. Residual + post_feedforward_layernorm
        current = try applyResidualAdd(input: current, add: mlpOut, seqLen: numPatches, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf)
        current = try applyRMSNorm(input: current, weight: weights.postFeedforwardLayernorm, seqLen: numPatches, hiddenSize: config.hiddenSize, cmdBuf: cmdBuf)

        return current
    }

    private func applyVisionAttention(input: MTLBuffer, weights: VisionLayerWeights, numPatches: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        // Q, K, V projections
        let q = try applyQuantizedMatmul(input: input, weights: weights.selfAttnQProj, seqLen: numPatches, output: qBuffer, cmdBuf: cmdBuf)
        let k = try applyQuantizedMatmul(input: input, weights: weights.selfAttnKProj, seqLen: numPatches, output: kBuffer, cmdBuf: cmdBuf)
        let v = try applyQuantizedMatmul(input: input, weights: weights.selfAttnVProj, seqLen: numPatches, output: vBuffer, cmdBuf: cmdBuf)

        // Q/K norm
        let qNormed = try applyHeadNorm(input: q, weight: weights.qNorm, seqLen: numPatches, numHeads: config.numAttentionHeads, headDim: config.headDim, cmdBuf: cmdBuf)
        let kNormed = try applyHeadNorm(input: k, weight: weights.kNorm, seqLen: numPatches, numHeads: config.numAttentionHeads, headDim: config.headDim, cmdBuf: cmdBuf)

        // Attention
        let attnOut = try applyAttention(q: qNormed, k: kNormed, v: v, numPatches: numPatches, numHeads: config.numAttentionHeads, headDim: config.headDim, output: attnOutBuffer, cmdBuf: cmdBuf)

        // O projection
        return try applyQuantizedMatmul(input: attnOut, weights: weights.selfAttnOProj, seqLen: numPatches, output: tempBuffer, cmdBuf: cmdBuf)
    }

    private func applyHeadNorm(input: MTLBuffer, weight: MTLBuffer, seqLen: Int, numHeads: Int, headDim: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let output = input

        let pso = try engine.pipeline(named: "vision_head_norm")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)

        enc.setBuffer(input, offset: 0, index: 0)
        enc.setBuffer(weight, offset: 0, index: 1)
        enc.setBuffer(output, offset: 0, index: 2)

        var numHeads_ = UInt32(numHeads)
        enc.setBytes(&numHeads_, length: 4, index: 3)
        var headDim_ = UInt32(headDim)
        enc.setBytes(&headDim_, length: 4, index: 4)
        var seqLen_ = UInt32(seqLen)
        enc.setBytes(&seqLen_, length: 4, index: 5)
        var eps = config.rmsNormEps
        enc.setBytes(&eps, length: 4, index: 6)

        let grid = MTLSize(width: numHeads * headDim, height: seqLen, depth: 1)
        let tg = engine.threadgroupSize2D(pso, grid: (numHeads * headDim, seqLen))
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()

        return output
    }

    private func applyAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, numPatches: Int, numHeads: Int, headDim: Int, output: MTLBuffer, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let pso = try engine.pipeline(named: "vision_attention")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)

        enc.setBuffer(q, offset: 0, index: 0)
        enc.setBuffer(k, offset: 0, index: 1)
        enc.setBuffer(v, offset: 0, index: 2)
        enc.setBuffer(output, offset: 0, index: 3)

        var numPatches_ = UInt32(numPatches)
        enc.setBytes(&numPatches_, length: 4, index: 4)
        var numHeads_ = UInt32(numHeads)
        enc.setBytes(&numHeads_, length: 4, index: 5)
        var headDim_ = UInt32(headDim)
        enc.setBytes(&headDim_, length: 4, index: 6)

        let grid = MTLSize(width: numHeads * headDim, height: numPatches, depth: 1)
        let tg = engine.threadgroupSize2D(pso, grid: (numHeads * headDim, numPatches))
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()

        return output
    }

    private func applyVisionMLP(input: MTLBuffer, weights: VisionLayerWeights, numPatches: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        // Gate projection: [numPatches, 768] -> [numPatches, 3072]
        let gate = try applyQuantizedMatmul(input: input, weights: weights.mlpGateProj, seqLen: numPatches, output: mlpBuffer, cmdBuf: cmdBuf)

        // Up projection: [numPatches, 768] -> [numPatches, 3072]
        let up = try applyQuantizedMatmul(input: input, weights: weights.mlpUpProj, seqLen: numPatches, output: tempBuffer, cmdBuf: cmdBuf)

        // SiLU(gate) * up
        let gated = try applyGateMultiply(gate: gate, up: up, count: numPatches * config.intermediateSize, cmdBuf: cmdBuf)

        // Down projection: [numPatches, 3072] -> [numPatches, 768]
        return try applyQuantizedMatmul(input: gated, weights: weights.mlpDownProj, seqLen: numPatches, output: tempBuffer, cmdBuf: cmdBuf)
    }

    private func applyGateMultiply(gate: MTLBuffer, up: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let output = mlpBuffer

        let pso = try engine.pipeline(named: "vision_gate_multiply")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)

        enc.setBuffer(gate, offset: 0, index: 0)
        enc.setBuffer(up, offset: 0, index: 1)
        enc.setBuffer(output, offset: 0, index: 2)

        var count_ = UInt32(count)
        enc.setBytes(&count_, length: 4, index: 3)

        let grid = MTLSize(width: count, height: 1, depth: 1)
        let tg = engine.threadgroupSize1D(pso, count: count)
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()

        return output
    }

    // ── Utility kernels ───────────────────────────────

    private func applyRMSNorm(input: MTLBuffer, weight: MTLBuffer, seqLen: Int, hiddenSize: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let output = tempBuffer

        let pso = try engine.pipeline(named: "rms_norm_seq")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)

        enc.setBuffer(input, offset: 0, index: 0)
        enc.setBuffer(weight, offset: 0, index: 1)
        enc.setBuffer(output, offset: 0, index: 2)

        var N = UInt32(hiddenSize)
        enc.setBytes(&N, length: 4, index: 3)
        var eps = config.rmsNormEps
        enc.setBytes(&eps, length: 4, index: 4)
        var sl = UInt32(seqLen)
        enc.setBytes(&sl, length: 4, index: 5)

        let grid = MTLSize(width: hiddenSize, height: seqLen, depth: 1)
        let tg = engine.threadgroupSize2D(pso, grid: (hiddenSize, seqLen))
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()

        return output
    }

    private func applyResidualAdd(input: MTLBuffer, add: MTLBuffer, seqLen: Int, hiddenSize: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let output = residualBuffer

        let pso = try engine.pipeline(named: "vision_residual_add")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)

        enc.setBuffer(input, offset: 0, index: 0)
        enc.setBuffer(add, offset: 0, index: 1)
        enc.setBuffer(output, offset: 0, index: 2)

        var count = UInt32(seqLen * hiddenSize)
        enc.setBytes(&count, length: 4, index: 3)

        let grid = MTLSize(width: seqLen * hiddenSize, height: 1, depth: 1)
        let tg = engine.threadgroupSize1D(pso, count: seqLen * hiddenSize)
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()

        return output
    }

    private func applyEmbeddingProjection(input: MTLBuffer, numPatches: Int, output: MTLBuffer, cmdBuf: MTLCommandBuffer) throws {
        let pso = try engine.pipeline(named: "vision_embedding_projection_quantized")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)

        enc.setBuffer(input, offset: 0, index: 0)
        enc.setBuffer(weights.embeddingProjectionWeight, offset: 0, index: 1)
        enc.setBuffer(weights.embeddingProjectionScales, offset: 0, index: 2)
        enc.setBuffer(weights.embeddingProjectionBiases, offset: 0, index: 3)
        enc.setBuffer(output, offset: 0, index: 4)

        var inFeatures = UInt32(768)   // Vision hidden size
        enc.setBytes(&inFeatures, length: 4, index: 5)
        var outFeatures = UInt32(2560)  // Text hidden size
        enc.setBytes(&outFeatures, length: 4, index: 6)
        var np = UInt32(numPatches)
        enc.setBytes(&np, length: 4, index: 7)
        var packedSize = UInt32(96)    // 768 / 8
        enc.setBytes(&packedSize, length: 4, index: 8)
        var groupSize = UInt32(64)
        enc.setBytes(&groupSize, length: 4, index: 9)
        var numGroups = UInt32(12)     // 768 / 64
        enc.setBytes(&numGroups, length: 4, index: 10)

        let grid = MTLSize(width: 2560, height: numPatches, depth: 1)
        let tg = engine.threadgroupSize2D(pso, grid: (2560, numPatches))
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()
    }
}