import Foundation
import Metal
import MetalPerformanceShaders
import Accelerate

/// EmbeddingGemma configuration
public struct EmbeddingGemmaConfig: Codable {
    public let hiddenSize: Int
    public let numHiddenLayers: Int
    public let vocabSize: Int
    public let numAttentionHeads: Int
    public let numKeyValueHeads: Int
    public let headDim: Int
    public let intermediateSize: Int
    public let maxPositionEmbeddings: Int
    public let slidingWindow: Int
    public let rmsNormEps: Float
    public let ropeTheta: Float
    public let useBidirectionalAttention: Bool
    public let layerTypes: [String]

    enum CodingKeys: String, CodingKey {
        case hiddenSize = "hidden_size"
        case numHiddenLayers = "num_hidden_layers"
        case vocabSize = "vocab_size"
        case numAttentionHeads = "num_attention_heads"
        case numKeyValueHeads = "num_key_value_heads"
        case headDim = "head_dim"
        case intermediateSize = "intermediate_size"
        case maxPositionEmbeddings = "max_position_embeddings"
        case slidingWindow = "sliding_window"
        case rmsNormEps = "rms_norm_eps"
        case ropeTheta = "rope_theta"
        case useBidirectionalAttention = "use_bidirectional_attention"
        case layerTypes = "layer_types"
    }

    public static func load(from modelDir: String) throws -> Self {
        let url = URL(fileURLWithPath: modelDir).appendingPathComponent("config.json")
        let data = try Data(contentsOf: url)
        return try JSONDecoder().decode(Self.self, from: data)
    }
}

public final class EmbeddingGemmaModel: @unchecked Sendable {
    public let config: EmbeddingGemmaConfig
    public let engine: MarkBaseEngine
    public let tokenizer: Tokenizer
    public let reader: SafeTensorsReader

    public var embedTokens: MTLBuffer!
    public var finalNorm: MTLBuffer!
    public var layerNorms: [[MTLBuffer]] = []
    public var qProjs: [MTLBuffer] = []
    public var kProjs: [MTLBuffer] = []
    public var vProjs: [MTLBuffer] = []
    public var oProjs: [MTLBuffer] = []
    public var qNorms: [MTLBuffer] = []
    public var kNorms: [MTLBuffer] = []
    public var gateProjs: [MTLBuffer] = []
    public var upProjs: [MTLBuffer] = []
    public var downProjs: [MTLBuffer] = []

    public init(modelDir: String, engine: MarkBaseEngine) throws {
        self.engine = engine
        self.config = try EmbeddingGemmaConfig.load(from: modelDir)
        self.tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        self.reader = try SafeTensorsReader(path: modelDir + "/model.safetensors")
        try loadWeights()
        print("✓ EmbeddingGemma loaded (\(config.numHiddenLayers) layers, hidden=\(config.hiddenSize))")
    }

    private func loadWeights() throws {
        let embedData = try readTensor("embed_tokens.weight")
        embedTokens = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
        for i in 0..<config.numHiddenLayers {
            let p = "layers.\(i)"
            layerNorms.append([
                try loadBuffer("\(p).input_layernorm.weight"),
                try loadBuffer("\(p).pre_feedforward_layernorm.weight"),
                try loadBuffer("\(p).post_attention_layernorm.weight"),
                try loadBuffer("\(p).post_feedforward_layernorm.weight"),
            ])
            qProjs.append(try loadAndTranspose("\(p).self_attn.q_proj.weight", rows: config.hiddenSize, cols: config.hiddenSize))
            kProjs.append(try loadAndTranspose("\(p).self_attn.k_proj.weight", rows: config.numKeyValueHeads * config.headDim, cols: config.hiddenSize))
            vProjs.append(try loadAndTranspose("\(p).self_attn.v_proj.weight", rows: config.numKeyValueHeads * config.headDim, cols: config.hiddenSize))
            oProjs.append(try loadAndTranspose("\(p).self_attn.o_proj.weight", rows: config.numAttentionHeads * config.headDim, cols: config.hiddenSize))
            qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight"))
            kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight"))
            gateProjs.append(try loadAndTranspose("\(p).mlp.gate_proj.weight", rows: config.intermediateSize, cols: config.hiddenSize))
            upProjs.append(try loadAndTranspose("\(p).mlp.up_proj.weight", rows: config.intermediateSize, cols: config.hiddenSize))
            downProjs.append(try loadAndTranspose("\(p).mlp.down_proj.weight", rows: config.hiddenSize, cols: config.intermediateSize))
        }
        let fnData = try readTensor("norm.weight")
        finalNorm = engine.device.makeBuffer(bytes: fnData, length: fnData.count * 4)!
    }

    public func embed(text: String, maxLen: Int = 2048) throws -> [Float] {
        var tokens = tokenizer.encode(text: text)
        if tokens.count > maxLen { tokens = Array(tokens.prefix(maxLen)) }
        guard !tokens.isEmpty else { return [] }

        let seqLen = tokens.count, hs = config.hiddenSize

        // Test 1: Embedding lookup only
        print("  TEST: Embedding lookup...")
        let cmdBuf1 = engine.commandQueue.makeCommandBuffer()!
        let inputBuf = try lookupEmbeddings(tokens: tokens, cmdBuf: cmdBuf1)
        cmdBuf1.commit()
        cmdBuf1.waitUntilCompleted()
        print("  TEST: Embedding lookup OK")

        // Test 2: Single layer forward (layer 0 only)
        print("  TEST: Layer 0 forward...")
        let cmdBuf2 = engine.commandQueue.makeCommandBuffer()!
        var hidden = try forwardLayerDebug(hidden: inputBuf, layerIdx: 0, seqLen: seqLen, cmdBuf: cmdBuf2)
        cmdBuf2.commit()
        cmdBuf2.waitUntilCompleted()
        print("  TEST: Layer 0 OK")

        // Full forward pass
        print("  TEST: Full forward pass...")
        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
        hidden = inputBuf
        for idx in 0..<config.numHiddenLayers {
            if idx % 6 == 0 { print("  Layer \(idx)/\(config.numHiddenLayers)") }
            hidden = try forwardLayerDebug(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
        }
        print("  TEST: All layers OK")

        let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs, cmdBuf: cmdBuf)
        cmdBuf.commit()
        cmdBuf.waitUntilCompleted()

        let data = engine.readFloats(from: output, count: seqLen * hs)
        var embedding = [Float](repeating: 0, count: hs)
        for i in 0..<seqLen {
            let start = i * hs
            for j in 0..<hs { embedding[j] += data[start + j] }
        }
        let n = Float(seqLen)
        for i in 0..<hs { embedding[i] /= n }
        var norm: Float = 0
        for i in 0..<hs { norm += embedding[i] * embedding[i] }
        norm = sqrt(norm)
        if norm > 0 { for i in 0..<hs { embedding[i] /= norm } }
        return embedding
    }

    private func readTensor(_ name: String) throws -> [Float] {
        guard let desc = reader.tensor(named: name) else { throw WeightError.tensorNotFound(name) }
        let data = try reader.read(tensor: desc)
        switch desc.dtype {
        case .f32: return data.withUnsafeBytes { Array(UnsafeBufferPointer(start: $0.baseAddress?.assumingMemoryBound(to: Float.self), count: data.count/4)) }
        case .bf16: return try SafeTensorsReader.bf16ToFloat32(data)
        default: throw WeightError.unsupportedDtype(desc.dtype.rawValue)
        }
    }

    private func loadBuffer(_ name: String) throws -> MTLBuffer {
        let data = try readTensor(name)
        return engine.device.makeBuffer(bytes: data, length: data.count * 4)!
    }

    private func loadAndTranspose(_ name: String, rows: Int, cols: Int) throws -> MTLBuffer {
        // Load [rows, cols] and transpose to [cols, rows] for MPS matmul C = A × B
        let data = try readTensor(name)
        var transposed = [Float](repeating: 0, count: data.count)
        for r in 0..<rows {
            for c in 0..<cols {
                transposed[c * rows + r] = data[r * cols + c]
            }
        }
        return engine.device.makeBuffer(bytes: transposed, length: transposed.count * 4)!
    }

    private func lookupEmbeddings(tokens: [Int], cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let seqLen = tokens.count, hs = config.hiddenSize
        print("    lookupEmbeddings: seqLen=\(seqLen), hs=\(hs)")
        // Read embedding table to CPU
        let embedPtr = embedTokens.contents().assumingMemoryBound(to: Float.self)
        let embedCount = embedTokens.length / 4
        let embedArray = Array(UnsafeBufferPointer(start: embedPtr, count: embedCount))
        print("    Read \(embedCount) floats from embedTokens")
        
        // Lookup embeddings for tokens
        var embedData = [Float](repeating: 0, count: seqLen * hs)
        for (i, token) in tokens.enumerated() {
            let start = i * hs
            let srcStart = token * hs
            embedData[start..<start+hs] = embedArray[srcStart..<srcStart+hs]
        }
        print("    Looked up \(seqLen) tokens")
        
        let buf = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
        print("    Created MTLBuffer")
        return buf
    }

    private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let output = engine.device.makeBuffer(length: count * 4)!
        let enc = cmdBuf.makeComputeCommandEncoder()!
        defer { enc.endEncoding() }
        let pso = try engine.pipeline(named: "rms_norm")
        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 c = UInt32(count), e: Float = config.rmsNormEps
        enc.setBytes(&c, length: 4, index: 3)
        enc.setBytes(&e, length: 4, index: 4)
        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
        return output
    }

    private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
        // Use MPS for optimized matrix multiplication on Apple Silicon
        // Weight is stored transposed [k, n] for C = A × B
        let descA = MPSMatrixDescriptor(rows: m, columns: k, rowBytes: k * 4, dataType: .float32)
        let descB = MPSMatrixDescriptor(rows: k, columns: n, rowBytes: n * 4, dataType: .float32)
        let descC = MPSMatrixDescriptor(rows: m, columns: n, rowBytes: n * 4, dataType: .float32)
        let matA = MPSMatrix(buffer: input, descriptor: descA)
        let matB = MPSMatrix(buffer: weight, descriptor: descB)
        let matC = MPSMatrix(buffer: output, descriptor: descC)
        
        let matMul = MPSMatrixMultiplication(device: engine.device,
                                              transposeLeft: false,
                                              transposeRight: false,
                                              resultRows: m,
                                              resultColumns: n,
                                              interiorColumns: k,
                                              alpha: 1.0,
                                              beta: 0.0)
        matMul.encode(commandBuffer: cmdBuf, leftMatrix: matA, rightMatrix: matB, resultMatrix: matC)
    }

    private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
        let enc = cmdBuf.makeComputeCommandEncoder()!
        defer { enc.endEncoding() }
        let pso = try engine.pipeline(named: "eltwise_add")
        enc.setComputePipelineState(pso)
        enc.setBuffer(a, offset: 0, index: 0)
        enc.setBuffer(b, offset: 0, index: 1)
        enc.setBuffer(output, offset: 0, index: 2)
        var c = UInt32(count)
        enc.setBytes(&c, length: 4, index: 3)
        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
    }

    private func geluMul(gate: MTLBuffer, up: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
        let enc = cmdBuf.makeComputeCommandEncoder()!
        defer { enc.endEncoding() }
        let pso = try engine.pipeline(named: "gelu_mul_kernel")
        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 c = UInt32(count)
        enc.setBytes(&c, length: 4, index: 3)
        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
    }

    private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, cmdBuf: MTLCommandBuffer) throws {
        let enc = cmdBuf.makeComputeCommandEncoder()!
        defer { enc.endEncoding() }
        let pso = try engine.pipeline(named: "apply_rope")
        enc.setComputePipelineState(pso)
        enc.setBuffer(q, offset: 0, index: 0)
        enc.setBuffer(k, offset: 0, index: 1)
        var sl = UInt32(seqLen), hd = UInt32(headDim), nh = UInt32(numHeads)
        var rt: Float = Float(config.ropeTheta)
        enc.setBytes(&sl, length: 4, index: 2)
        enc.setBytes(&hd, length: 4, index: 3)
        enc.setBytes(&nh, length: 4, index: 4)
        enc.setBytes(&rt, length: 4, index: 5)
        enc.dispatchThreads(MTLSize(width: numHeads * headDim / 2, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, numHeads * headDim / 2), height: 1, depth: 1))
    }

    private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
        let enc = cmdBuf.makeComputeCommandEncoder()!
        defer { enc.endEncoding() }
        let pso = try engine.pipeline(named: "bidirectional_sliding_attn")
        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 sl = UInt32(seqLen), hd = UInt32(config.headDim), nh = UInt32(config.numAttentionHeads)
        var nkv = UInt32(config.numKeyValueHeads), sw = UInt32(config.slidingWindow)
        var scale: Float = 1.0 / sqrt(Float(config.headDim))
        enc.setBytes(&sl, length: 4, index: 4)
        enc.setBytes(&hd, length: 4, index: 5)
        enc.setBytes(&nh, length: 4, index: 6)
        enc.setBytes(&nkv, length: 4, index: 7)
        enc.setBytes(&sw, length: 4, index: 8)
        enc.setBytes(&scale, length: 4, index: 9)
        let tgMem = config.slidingWindow * 4
        enc.setThreadgroupMemoryLength(tgMem, index: 0)
        enc.dispatchThreads(MTLSize(width: seqLen * config.numAttentionHeads, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, seqLen * config.numAttentionHeads), height: 1, depth: 1))
    }

    private func forwardLayerDebug(hidden: MTLBuffer, layerIdx: Int, seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
        let hs = config.hiddenSize, device = engine.device
        let hDim = config.headDim, nH = config.numAttentionHeads, nKV = config.numKeyValueHeads
        let intermedi = config.intermediateSize

        print("    Residual copy...")
        let resid = device.makeBuffer(length: seqLen * hs * 4)!
        let blit = cmdBuf.makeBlitCommandEncoder()!
        blit.copy(from: hidden, sourceOffset: 0, to: resid, destinationOffset: 0, size: seqLen * hs * 4)
        blit.endEncoding()

        print("    Input norm...")
        let h1 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][0], count: seqLen * hs, cmdBuf: cmdBuf)

        print("    QKV projections...")
        let qBuf = device.makeBuffer(length: seqLen * nH * hDim * 4)!
        let kBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
        let vBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
        try matmulSeq(input: h1, weight: qProjs[layerIdx], output: qBuf, m: seqLen, k: hs, n: nH * hDim, cmdBuf: cmdBuf)
        print("    Q done")
        try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
        print("    K done")
        try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
        print("    V done")

        print("    RoPE...")
        try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, cmdBuf: cmdBuf)
        print("    RoPE done")

        print("    Attention...")
        let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
        try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)
        print("    Attention done")

        print("    O projection...")
        let h2 = device.makeBuffer(length: seqLen * hs * 4)!
        try matmulSeq(input: attnOut, weight: oProjs[layerIdx], output: h2, m: seqLen, k: nH * hDim, n: hs, cmdBuf: cmdBuf)
        print("    O done")

        print("    Post-attn norm...")
        let h2n = try applyRmsNorm(input: h2, weight: layerNorms[layerIdx][2], count: seqLen * hs, cmdBuf: cmdBuf)

        print("    Add residual 1...")
        try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)

        print("    Pre-FF norm...")
        let h3 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][1], count: seqLen * hs, cmdBuf: cmdBuf)

        print("    MLP gate/up...")
        let gate = device.makeBuffer(length: seqLen * intermedi * 4)!
        let up = device.makeBuffer(length: seqLen * intermedi * 4)!
        try matmulSeq(input: h3, weight: gateProjs[layerIdx], output: gate, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
        print("    Gate done")
        try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
        print("    Up done")

        print("    GELU mul...")
        let gated = device.makeBuffer(length: seqLen * intermedi * 4)!
        try geluMul(gate: gate, up: up, output: gated, count: seqLen * intermedi, cmdBuf: cmdBuf)
        print("    GELU done")

        print("    Down projection...")
        let h4 = device.makeBuffer(length: seqLen * hs * 4)!
        try matmulSeq(input: gated, weight: downProjs[layerIdx], output: h4, m: seqLen, k: intermedi, n: hs, cmdBuf: cmdBuf)
        print("    Down done")

        print("    Post-FF norm...")
        let h4n = try applyRmsNorm(input: h4, weight: layerNorms[layerIdx][3], count: seqLen * hs, cmdBuf: cmdBuf)

        print("    Add residual 2...")
        try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
        print("    Layer \(layerIdx) complete")

        return hidden
    }
}