import XCTest
@testable import MarkBase

final class EmbeddingTest: XCTestCase {
    var engine: MarkBaseEngine!
    var model: E4BModel!
    var embeddingModel: TextEmbeddingModel!
    let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"

    override func setUp() {
        super.setUp()
        guard FileManager.default.fileExists(atPath: modelDir + "/model.safetensors") else {
            return
        }
        engine = try? MarkBaseEngine(autoCompile: true)
        model = try? E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        embeddingModel = try? TextEmbeddingModel(modelDir: modelDir, engine: engine, config: TextEmbeddingConfig())
    }

    func testEmbeddingDimension() throws {
        try XCTSkipIf(embeddingModel == nil, "Model not found")
        let embedding = try embeddingModel.embed(text: "Hello world")
        XCTAssertEqual(embedding.count, 2560, "Embedding dimension should be 2560")
    }

    func testEmbeddingNormalized() throws {
        try XCTSkipIf(embeddingModel == nil, "Model not found")
        let embedding = try embeddingModel.embed(text: "Test text")
        let norm = sqrt(embedding.reduce(0) { $0 + $1 * $1 })
        XCTAssertEqual(norm, 1.0, accuracy: 0.001, "Embedding should be L2 normalized")
    }

    func testSimilarSentences() throws {
        try XCTSkipIf(embeddingModel == nil, "Model not found")
        let e1 = try embeddingModel.embed(text: "The cat is sitting on the mat")
        let e2 = try embeddingModel.embed(text: "A cat rests on a rug")
        let e3 = try embeddingModel.embed(text: "The stock market crashed today")

        let sim12 = cosineSimilarity(e1, e2)
        let sim13 = cosineSimilarity(e1, e3)
        print("Similar(cat, cat): \(sim12)")
        print("Similar(cat, stock): \(sim13)")
        XCTAssertGreaterThan(sim12, sim13, "Similar sentences should have higher cosine similarity")
    }

    func testDifferentLengths() throws {
        try XCTSkipIf(embeddingModel == nil, "Model not found")
        let e1 = try embeddingModel.embed(text: "Hi")
        let e2 = try embeddingModel.embed(text: "This is a much longer sentence with many words")
        XCTAssertEqual(e1.count, e2.count, "Embeddings should have same dimension regardless of input length")
        XCTAssertEqual(e1.count, 2560)
    }

    func testEmptyInput() throws {
        try XCTSkipIf(embeddingModel == nil, "Model not found")
        let embedding = try embeddingModel.embed(text: "")
        XCTAssertEqual(embedding.count, 0, "Empty input should return empty embedding")
    }

    func testBatchEmbedding() throws {
        try XCTSkipIf(embeddingModel == nil, "Model not found")
        let texts = ["Hello", "World", "Test"]
        let embeddings = try embeddingModel.embedBatch(texts: texts)
        XCTAssertEqual(embeddings.count, 3)
        for embedding in embeddings {
            XCTAssertEqual(embedding.count, 2560)
        }
    }

    private func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float {
        guard a.count == b.count, !a.isEmpty else { return 0 }
        var dot: Float = 0, normA: Float = 0, normB: Float = 0
        for i in 0..<a.count {
            dot += a[i] * b[i]
            normA += a[i] * a[i]
            normB += b[i] * b[i]
        }
        return dot / (sqrt(normA) * sqrt(normB))
    }
}