"""Tests for CPUEmbeddingEngine."""

from __future__ import annotations

import sys
import types

import pytest

from sagellm_core.engines import CPUEmbeddingEngine, get_embedding_engine_class


class _FakeEmbeddingResult:
    def __init__(self, vectors: list[list[float]]) -> None:
        self._vectors = vectors

    def tolist(self) -> list[list[float]]:
        return self._vectors


class _FakeSentenceTransformer:
    def __init__(self, model_id: str, device: str = "cpu") -> None:
        self.model_id = model_id
        self.device = device

    def encode(self, texts: list[str], **_: object) -> _FakeEmbeddingResult:
        return _FakeEmbeddingResult(
            [[float(len(text)), float(index)] for index, text in enumerate(texts)]
        )


@pytest.fixture
def fake_sentence_transformers(monkeypatch: pytest.MonkeyPatch) -> None:
    module = types.SimpleNamespace(SentenceTransformer=_FakeSentenceTransformer)
    monkeypatch.setitem(sys.modules, "sentence_transformers", module)


@pytest.mark.asyncio
async def test_cpu_embedding_engine_encode_shape(fake_sentence_transformers: None) -> None:
    engine = CPUEmbeddingEngine(model_id="sentence-transformers/all-MiniLM-L6-v2", device="cpu")

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

    assert len(vectors) == 3
    assert all(len(vector) == 2 for vector in vectors)
    assert vectors[0][0] == 5.0

    await engine.stop()


@pytest.mark.asyncio
async def test_cpu_embedding_engine_health_check(fake_sentence_transformers: None) -> None:
    engine = CPUEmbeddingEngine(model_id="sentence-transformers/all-MiniLM-L6-v2")

    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


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