"""Unit tests for coalex.ext.embedding."""

from __future__ import annotations

import json

from coalex.ext.embedding import EmbeddingSpan, embedding_span


def test_embedding_span_decorator_sets_kind(memory_exporter):
    @embedding_span(model_name="text-embedding-3-small")
    def embed(text: str) -> list[float]:
        return [0.1, 0.2, 0.3]

    embed(text="hello world")
    spans = memory_exporter.get_finished_spans()
    assert len(spans) == 1
    assert spans[0].attributes["openinference.span.kind"] == "EMBEDDING"


def test_embedding_span_decorator_records_model_name(memory_exporter):
    @embedding_span(model_name="text-embedding-ada-002")
    def embed(text: str) -> list[float]:
        return [0.1]

    embed(text="test")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["embedding.model_name"] == "text-embedding-ada-002"


def test_embedding_span_decorator_records_vector(memory_exporter):
    @embedding_span(model_name="m")
    def embed(text: str) -> list[float]:
        return [0.5, 0.6]

    embed(text="encode me")
    spans = memory_exporter.get_finished_spans()
    encoded = json.loads(spans[0].attributes["embedding.embeddings"])
    assert encoded[0]["embedding.vector"] == [0.5, 0.6]
    assert encoded[0]["embedding.text"] == "encode me"


def test_embedding_span_decorator_records_input_text(memory_exporter):
    @embedding_span()
    def embed(text: str) -> list[float]:
        return [0.1]

    embed(text="my text")
    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["input.value"] == "my text"


def test_embedding_span_context_manager(memory_exporter):
    with EmbeddingSpan("embed_step", model_name="text-embedding-3-large") as e:
        e.set_text("hello")
        e.set_vector([0.1, 0.2], text="hello")

    spans = memory_exporter.get_finished_spans()
    assert spans[0].attributes["openinference.span.kind"] == "EMBEDDING"
    assert spans[0].attributes["embedding.model_name"] == "text-embedding-3-large"
    assert spans[0].attributes["input.value"] == "hello"
    encoded = json.loads(spans[0].attributes["embedding.embeddings"])
    assert encoded[0]["embedding.vector"] == [0.1, 0.2]


# ---------------------------------------------------------------------------
# Async decorator
# ---------------------------------------------------------------------------


async def test_embedding_span_async_records_vector(memory_exporter):
    @embedding_span(model_name="text-embedding-3-small")
    async def embed(text: str) -> list[float]:
        return [0.3, 0.4, 0.5]

    await embed(text="async text")
    spans = memory_exporter.get_finished_spans()
    assert len(spans) == 1
    assert spans[0].attributes["openinference.span.kind"] == "EMBEDDING"
    assert spans[0].attributes["embedding.model_name"] == "text-embedding-3-small"
    assert spans[0].attributes["input.value"] == "async text"
    encoded = json.loads(spans[0].attributes["embedding.embeddings"])
    assert encoded[0]["embedding.vector"] == [0.3, 0.4, 0.5]
