import Foundation
import MLX
import MLXEmbedders
import MLXLMCommon

/// Everything that runs inside `ModelContainer.perform { ... }`.
/// Pure function: given model/tokenizer/pool + texts, returns `[[Float]]`.
enum Batching {
    static func embed(
        model: EmbeddingModel,
        tokenizer: any Tokenizer,
        pooling: Pooling,
        texts: [String],
        kind: InputKind,
        maxLen: Int
    ) -> [[Float]] {
        precondition(!texts.isEmpty, "embed called with empty batch")

        // 1. Prompt wrap + tokenize
        let wrapped = texts.map { Prompt.wrap($0, kind) }
        var rows: [[Int32]] = wrapped.map { text in
            let ids = tokenizer.encode(text: text, addSpecialTokens: true)
            // Truncate to maxLen hard cap (rare; buckets usually govern this).
            let capped = ids.count > maxLen ? Array(ids.prefix(maxLen)) : ids
            return capped.map { Int32($0) }
        }

        // 2. Pick length bucket from the longest row.
        let maxRow = rows.map { $0.count }.max() ?? 1
        let length = Buckets.length(for: max(maxRow, 1), maxLen: maxLen)

        // 3. Pad rows to length; build attention mask.
        let padId: Int32 = {
            if let eos = tokenizer.eosTokenId { return Int32(eos) }
            return 0
        }()

        let batchSize = rows.count
        var idsFlat = [Int32](repeating: padId, count: batchSize * length)
        var maskFlat = [Int32](repeating: 0, count: batchSize * length)
        for (i, row) in rows.enumerated() {
            let n = min(row.count, length)
            let base = i * length
            for j in 0..<n {
                idsFlat[base + j] = row[j]
                maskFlat[base + j] = 1
            }
        }
        rows.removeAll(keepingCapacity: false)

        let inputIds = MLXArray(idsFlat, [batchSize, length])
        let attentionMask = MLXArray(maskFlat, [batchSize, length])

        // 4. Forward.
        let out = model(
            inputIds,
            positionIds: nil,
            tokenTypeIds: nil,
            attentionMask: attentionMask
        )

        // 5. Last-token pool + L2 normalize. Pooling mask consumed as-is; it
        //    casts to int32 internally (see Pooling.swift: `_mask.asType(.int32)`).
        let pooled = pooling(out, mask: attentionMask, normalize: true)

        // 6. Evaluate before leaving the perform closure (MLXArray is not Sendable).
        eval(pooled)

        // 7. Copy to host memory as `[[Float]]`.
        let dim = pooled.dim(1)
        let host: [Float] = pooled.asArray(Float.self)
        var result = [[Float]]()
        result.reserveCapacity(batchSize)
        for i in 0..<batchSize {
            let base = i * dim
            result.append(Array(host[base..<(base + dim)]))
        }
        return result
    }
}