"""Embedding span decorator and context manager."""

from __future__ import annotations

import asyncio
import functools
import json
from collections.abc import Callable
from typing import Any, TypeVar

from coalex.ext._base import _get_tracer

F = TypeVar("F", bound=Callable[..., Any])


def embedding_span(
    name: str | None = None,
    model_name: str | None = None,
    query_arg: str = "text",
) -> Callable[[F], F]:
    """Decorator that wraps an embedding function with an OpenInference EMBEDDING span.

    Supports both sync and async functions. The span captures:
    - ``openinference.span.kind`` = ``EMBEDDING``
    - ``embedding.model_name`` — if provided
    - ``input.value`` — text being embedded
    - ``embedding.embeddings`` — JSON-encoded vector if function returns ``list[float]``

    Args:
        name: Span name. Defaults to the function name.
        model_name: Embedding model name to record.
        query_arg: Name of the kwarg containing the text to embed.
    """

    def decorator(fn: F) -> F:
        span_name = name or fn.__name__

        if asyncio.iscoroutinefunction(fn):

            @functools.wraps(fn)
            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
                tracer = _get_tracer()
                with tracer.start_as_current_span(span_name) as span:
                    span.set_attribute("openinference.span.kind", "EMBEDDING")
                    if model_name:
                        span.set_attribute("embedding.model_name", model_name)
                    text = kwargs.get(query_arg) or next((a for a in args if isinstance(a, str)), None)
                    if text:
                        span.set_attribute("input.value", str(text))
                    result = await fn(*args, **kwargs)
                    if isinstance(result, list) and result and isinstance(result[0], float):
                        encoded = json.dumps([{"embedding.vector": result, "embedding.text": str(text or "")}])
                        span.set_attribute("embedding.embeddings", encoded)
                    return result

            return async_wrapper  # type: ignore[return-value]

        @functools.wraps(fn)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            tracer = _get_tracer()
            with tracer.start_as_current_span(span_name) as span:
                span.set_attribute("openinference.span.kind", "EMBEDDING")
                if model_name:
                    span.set_attribute("embedding.model_name", model_name)
                text = kwargs.get(query_arg) or next((a for a in args if isinstance(a, str)), None)
                if text:
                    span.set_attribute("input.value", str(text))
                result = fn(*args, **kwargs)
                if isinstance(result, list) and result and isinstance(result[0], float):
                    encoded = json.dumps([{"embedding.vector": result, "embedding.text": str(text or "")}])
                    span.set_attribute("embedding.embeddings", encoded)
                return result

        return wrapper  # type: ignore[return-value]

    return decorator


class EmbeddingSpan:
    """Context manager for instrumenting an embedding call as an EMBEDDING span.

    Example::

        with EmbeddingSpan("embed", model_name="text-embedding-3-small") as e:
            e.set_text(text)
            vector = embed(text)
            e.set_vector(vector, text)
    """

    def __init__(self, name: str = "embedding", model_name: str | None = None) -> None:
        self._name = name
        self._model_name = model_name
        self._span: Any = None
        self._ctx: Any = None

    def __enter__(self) -> EmbeddingSpan:
        tracer = _get_tracer()
        self._ctx = tracer.start_as_current_span(self._name)
        self._span = self._ctx.__enter__()
        self._span.set_attribute("openinference.span.kind", "EMBEDDING")
        if self._model_name:
            self._span.set_attribute("embedding.model_name", self._model_name)
        return self

    def set_text(self, text: str) -> None:
        self._span.set_attribute("input.value", text)

    def set_vector(self, vector: list[float], text: str = "") -> None:
        encoded = json.dumps([{"embedding.vector": vector, "embedding.text": text}])
        self._span.set_attribute("embedding.embeddings", encoded)

    def __exit__(self, *exc_info: Any) -> bool | None:
        return self._ctx.__exit__(*exc_info)
