import Foundation
import Metal
import Accelerate

/// EmbeddingGemmaConfig - Configuration for EmbeddingGemma model
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)
    }
}

/// EmbeddingGemma - Google's 300M parameter embedding model
public final class EmbeddingGemmaModel: @unchecked Sendable {
    public let config: EmbeddingGemmaConfig
    public let engine: MarkBaseEngine
    public let tokenizer: Tokenizer
    public let reader: SafeTensorsReader

    // GPU Buffers
    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 hs = config.hiddenSize
        let intermedi = config.intermediateSize
        let nKV = config.numKeyValueHeads
        let hDim = config.headDim

        // Embedding table [vocab, hidden]
        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 loadBuffer("\(p).self_attn.q_proj.weight"))     // [hs, hs]
            kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight"))     // [nKV*hDim, hs]
            vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight"))     // [nKV*hDim, hs]
            oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight"))     // [hs, nH*hDim]
            qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight"))     // [hDim]
            kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight"))     // [hDim]
            gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight"))     // [intermedi, hs]
            upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight"))         // [intermedi, hs]
            downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight"))     // [hs, intermedi]
        }

        let fnData = try readTensor("norm.weight")
        finalNorm = engine.device.makeBuffer(bytes: fnData, length: fnData.count * 4)!
    }

    /// Generate embedding for text
    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

        // Embedding lookup
        let inputBuf = try lookupEmbeddings(tokens: tokens)

        // Forward through layers
        var hidden = inputBuf
        for idx in 0..<config.numHiddenLayers {
            hidden = try forwardLayer(hidden: hidden, layerIdx: idx, seqLen: seqLen)
        }

        // Final norm
        let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs)

        // Readback
        let data = engine.readFloats(from: output, count: seqLen * hs)

        // Mean pool + L2 normalize
        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
    }

    // MARK: - Helpers

    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 lookupEmbeddings(tokens: [Int]) throws -> MTLBuffer {
        let seqLen = tokens.count, hs = config.hiddenSize
        let buf = engine.device.makeBuffer(length: seqLen * hs * 4)!
        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
        let enc = cmdBuf.makeComputeCommandEncoder()!
        let pso = try engine.pipeline(named: "lookup_embeddings")
        enc.setComputePipelineState(pso)
        enc.setBuffer(embedTokens, offset: 0, index: 0)
        enc.setBytes(tokens, length: seqLen * MemoryLayout<Int>.size, index: 1)
        enc.setBuffer(buf, offset: 0, index: 2)
        var h = UInt32(hs), s = UInt32(seqLen), v = UInt32(config.vocabSize)
        enc.setBytes(&h, length: 4, index: 3)
        enc.setBytes(&s, length: 4, index: 4)
        enc.setBytes(&v, length: 4, index: 5)
        enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1),
                           threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
        enc.endEncoding()
        cmdBuf.commit(); cmdBuf.waitUntilCompleted()
        return buf
    }

    private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int) throws -> MTLBuffer {
        let output = engine.device.makeBuffer(length: count * 4)!
        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
        let enc = cmdBuf.makeComputeCommandEncoder()!
        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))
        enc.endEncoding()
        cmdBuf.commit(); cmdBuf.waitUntilCompleted()
        return output
    }

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

        let cmdBuf = engine.commandQueue.makeCommandBuffer()!

        // Residual
        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()

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

        // Q, K, V projections (using optimized matmul)
        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)
        try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
        try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)

        // RoPE
        try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, numKVHeads: nKV, cmdBuf: cmdBuf)

        // Q/K Norm
        try applyQKNorm(q: qBuf, k: kBuf, qNorm: qNorms[layerIdx], kNorm: kNorms[layerIdx], seqLen: seqLen, headDim: hDim, numHeads: nH, numKVHeads: nKV, cmdBuf: cmdBuf)

        // Bidirectional sliding window attention
        let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
        try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)

        // 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)

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

        // Add residual: hidden = resid + h2n
        try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)

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

        // 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)
        try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)

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

        // 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)

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

        // Add residual: hidden = hidden + h4n
        try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)

        cmdBuf.commit(); cmdBuf.waitUntilCompleted()
        return hidden
    }

    // MARK: - Metal Kernels

    private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
        let enc = cmdBuf.makeComputeCommandEncoder()!
        let pso = try engine.pipeline(named: "matmul_f32")
        enc.setComputePipelineState(pso)
        // Note: matmul_f32 only handles M=1, so we loop over positions
        // For production: implement a proper batch matmul kernel
        for i in 0..<m {
            let inputOffset = i * k * 4
            let outputOffset = i * n * 4
            enc.setBuffer(input, offset: inputOffset, index: 0)
            enc.setBuffer(weight, offset: 0, index: 1)
            enc.setBuffer(output, offset: outputOffset, index: 2)
            var mm: UInt32 = 1, kk = UInt32(k), nn = UInt32(n)
            enc.setBytes(&mm, length: 4, index: 3)
            enc.setBytes(&kk, length: 4, index: 4)
            enc.setBytes(&nn, length: 4, index: 5)
            enc.dispatchThreads(MTLSize(width: n, height: 1, depth: 1),
                               threadsPerThreadgroup: MTLSize(width: min(256, n), height: 1, depth: 1))
        }
        enc.endEncoding()
    }

    private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, numKVHeads: Int, cmdBuf: MTLCommandBuffer) throws {
        let enc = cmdBuf.makeComputeCommandEncoder()!
        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))
        enc.endEncoding()
    }

    private func applyQKNorm(q: MTLBuffer, k: MTLBuffer, qNorm: MTLBuffer, kNorm: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, numKVHeads: Int, cmdBuf: MTLCommandBuffer) throws {
        // Apply RMSNorm per head
        // TODO: Implement per-head normalization
    }

    private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
        let enc = cmdBuf.makeComputeCommandEncoder()!
        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  // shared memory for scores
        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))
        enc.endEncoding()
    }

    private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
        let enc = cmdBuf.makeComputeCommandEncoder()!
        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))
        enc.endEncoding()
    }

    private func geluMul(gate: MTLBuffer, up: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
        let enc = cmdBuf.makeComputeCommandEncoder()!
        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))
        enc.endEncoding()
    }
}