"""Tests for CPUEmbeddingEngine."""

from __future__ import annotations

import os
import time

import pytest

from sagellm_core.engines import CPUEmbeddingEngine, get_embedding_engine_class


@pytest.fixture
def real_embedding_model_id() -> str:
    """Real embedding model id used in tests."""
    pytest.importorskip("sentence_transformers")
    return os.getenv("SAGELLM_TEST_EMBEDDING_MODEL", "sentence-transformers/all-MiniLM-L6-v2")


@pytest.mark.asyncio
async def test_cpu_embedding_engine_encode_shape(real_embedding_model_id: str) -> None:
    engine = CPUEmbeddingEngine(model_id=real_embedding_model_id, device="cpu")

    await engine.start()
    vectors = await engine.encode(["hello", "world", "sage"])

    assert len(vectors) == 3
    assert all(len(vector) > 0 for vector in vectors)
    assert all(isinstance(value, float) for value in vectors[0])

    await engine.stop()


@pytest.mark.asyncio
async def test_cpu_embedding_engine_health_check(real_embedding_model_id: str) -> None:
    engine = CPUEmbeddingEngine(model_id=real_embedding_model_id)

    assert await engine.health_check() is False
    await engine.start()
    assert await engine.health_check() is True
    await engine.stop()
    assert await engine.health_check() is False


@pytest.mark.asyncio
async def test_cpu_embedding_engine_latency_for_10_texts(real_embedding_model_id: str) -> None:
    """Encoding 10 texts should finish quickly on CPU."""
    engine = CPUEmbeddingEngine(model_id=real_embedding_model_id)
    await engine.start()

    texts = [f"text-{index}" for index in range(10)]
    started_at = time.perf_counter()
    vectors = await engine.encode(texts)
    elapsed_s = time.perf_counter() - started_at

    assert len(vectors) == 10
    assert elapsed_s < 5.0

    await engine.stop()


def test_embedding_registry_has_cpu_backend() -> None:
    engine_class = get_embedding_engine_class("cpu")
    assert engine_class is CPUEmbeddingEngine
