import Metal

// Multimodal inference pipeline for 12B
// Handles audio/image processing and integration with text model

public final class MultimodalInference {
    public let model: MultimodalModel
    public let engine: MarkBaseEngine

    // Temporary buffers for multimodal embeddings
    private let audioEmbedBuffer: MTLBuffer
    private let visionEmbedBuffer: MTLBuffer

    public init(model: MultimodalModel) throws {
        self.model = model
        self.engine = model.engine

        let device = engine.device
        let hiddenSize = model.textModel.hiddenSize

        audioEmbedBuffer = device.makeBuffer(length: 1024 * hiddenSize * 4)!
        visionEmbedBuffer = device.makeBuffer(length: 1024 * hiddenSize * 4)!
    }

    // Complete multimodal inference pipeline
    public func generate(
        textTokens: [Int],
        audioFeatures: [[Float]]? = nil,
        imagePatches: [Float]? = nil,
        numImagePatches: Int = 0,
        precomputedVisionEmbedding: MTLBuffer? = nil,
        maxTokens: Int = 50
    ) throws -> [Int] {
        print("\n═══════════════════════════════════════")
        print("    Multimodal Inference Pipeline")
        print("═══════════════════════════════════════\n")

        var fullTokens = textTokens
        let hiddenSize = model.textModel.hiddenSize
        var audioTokenCount = 0
        var imageTokenCount = 0

        // ── Step 1: Process audio ──
        if let audio = audioFeatures {
            print("Step 1: Processing audio...")
            print("  Audio frames: \(audio.count)")

            fullTokens.append(model.boaTokenId)
            audioTokenCount = audio.count
            for _ in 0..<audioTokenCount {
                fullTokens.append(model.audioTokenId)
            }
            fullTokens.append(model.eoaTokenId)

            if let tower = model.audioTower {
                let flatFeatures = audio.flatMap { $0 }
                let inputBuffer = engine.device.makeBuffer(
                    bytes: flatFeatures,
                    length: flatFeatures.count * MemoryLayout<Float>.stride
                )!
                try tower.forward(
                    inputBuffer: inputBuffer,
                    seqLen: audioTokenCount,
                    outputBuffer: audioEmbedBuffer
                )
                print("  ✓ Audio towers forward done")
            }
        }

        // ── Step 2: Process image ──
        if let precomputed = precomputedVisionEmbedding {
            // Pre-computed pooled embedding — single IMAGE token
            print("Step 2: Using precomputed vision embedding")
            fullTokens.append(model.boiTokenId)
            fullTokens.append(model.imageTokenId)
            fullTokens.append(model.eoiTokenId)
            imageTokenCount = 1
            // Copy the pooled embedding into visionEmbedBuffer
            let cmdBuf = engine.commandQueue.makeCommandBuffer()!
            let blit = cmdBuf.makeBlitCommandEncoder()!
            blit.copy(from: precomputed, sourceOffset: 0,
                      to: visionEmbedBuffer, destinationOffset: 0,
                      size: min(precomputed.length, visionEmbedBuffer.length))
            blit.endEncoding()
            cmdBuf.commit()
            cmdBuf.waitUntilCompleted()
        } else if let patches = imagePatches, numImagePatches > 0 {
            print("Step 2: Processing image...")
            print("  Image patches: \(numImagePatches)")

            fullTokens.append(model.boiTokenId)
            imageTokenCount = numImagePatches
            for _ in 0..<imageTokenCount {
                fullTokens.append(model.imageTokenId)
            }
            fullTokens.append(model.eoiTokenId)

            let inputBuffer = engine.device.makeBuffer(
                bytes: patches,
                length: patches.count * MemoryLayout<Float>.stride
            )!
            if let tower = model.visionTowerFull {
                try tower.forward(patchEmbeddings: inputBuffer, numPatches: imageTokenCount, outputBuffer: visionEmbedBuffer)
            } else if let tower = model.visionTower {
                try tower.forward(patchEmbeddings: inputBuffer, numPatches: imageTokenCount, outputBuffer: visionEmbedBuffer)
            }
            print("  ✓ Vision tower forward done")
        }

        // ── Step 3: Pre-fill prompt with injection ──
        print("\nStep 3: Pre-filling \(fullTokens.count) tokens...")

        var generated = fullTokens
        var audioIdx = 0
        var imageIdx = 0

        for pos in 0..<fullTokens.count {
            let tokenId = fullTokens[pos]

            if tokenId == model.audioTokenId, audioIdx < audioTokenCount {
                let offset = audioIdx * hiddenSize * MemoryLayout<Float>.stride
                _ = try model.textModel.forwardFromHidden(hiddenBuffer: audioEmbedBuffer, offset: offset, position: pos)
                audioIdx += 1
            } else if tokenId == model.imageTokenId, imageIdx < imageTokenCount {
                let offset = imageIdx * hiddenSize * MemoryLayout<Float>.stride
                _ = try model.textModel.forwardFromHidden(hiddenBuffer: visionEmbedBuffer, offset: offset, position: pos)
                imageIdx += 1
            } else {
                _ = try model.textModel.forward(tokenId: tokenId, position: pos)
            }
        }

        // ── Step 4: Auto-regressive generation ──
        print("Step 4: Generating \(maxTokens) tokens...")

        let sampler = Sampler()
        
        for _ in 0..<maxTokens {
            let logits = try model.textModel.forward(
                tokenId: generated.last ?? 0,
                position: generated.count - 1
            )

            // Use sampler with unused token filtering
            let nextToken = sampler.sample(
                logits: logits,
                temperature: 0.7,
                topK: 50,
                topP: 0.95,
                filterUnusedTokens: true
            )

            generated.append(nextToken)
        }

        let newTokens = generated.count - fullTokens.count
        print("  ✓ Generated \(newTokens) tokens")

        return generated
    }
}