"""Retrieval span decorator and context manager for RAG pipelines."""

from __future__ import annotations

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

from opentelemetry.trace import Status, StatusCode

from coalex.ext._base import _get_tracer

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


@dataclasses.dataclass
class Document:
    """A retrieved document with OpenInference-compatible fields."""

    content: str
    id: str | None = None
    score: float | None = None
    metadata: dict[str, Any] = dataclasses.field(default_factory=dict)

    def to_openinference_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {"document.content": self.content}
        if self.id is not None:
            d["document.id"] = self.id
        if self.score is not None:
            d["document.score"] = self.score
        if self.metadata:
            d["document.metadata"] = json.dumps(self.metadata)
        return d


def encode_documents(docs: list[Document]) -> str:
    """Serialize documents to OpenInference retrieval.documents JSON string."""
    return json.dumps([d.to_openinference_dict() for d in docs])


def _document_summary(docs: list[Document]) -> str:
    """Build a human-readable summary of retrieved documents for output.value."""
    summary = f"Retrieved {len(docs)} document(s)"
    ids = [d.id for d in docs if d.id]
    if ids:
        summary += f": [{', '.join(ids)}]"
    return summary


def retrieval_span(name: str | None = None, query_arg: str = "query") -> Callable[[F], F]:
    """Decorator that wraps a retrieval function with an OpenInference RETRIEVER span.

    Supports both sync and async functions. The decorated function must return
    ``list[Document]``. The span captures:
    - ``openinference.span.kind`` = ``RETRIEVER``
    - ``input.value`` — query string (from ``query_arg`` kwarg or first ``str`` positional)
    - ``retrieval.documents`` — JSON-encoded list of documents

    Args:
        name: Span name. Defaults to the function name.
        query_arg: Name of the kwarg containing the query string.
    """

    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", "RETRIEVER")
                    q = kwargs.get(query_arg) or next((a for a in args if isinstance(a, str)), None)
                    if q:
                        span.set_attribute("input.value", str(q))
                    result = await fn(*args, **kwargs)
                    if isinstance(result, list) and result and isinstance(result[0], Document):
                        span.set_attribute("retrieval.documents", encode_documents(result))
                        span.set_attribute("output.value", _document_summary(result))
                    span.set_status(Status(StatusCode.OK))
                    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", "RETRIEVER")
                q = kwargs.get(query_arg) or next((a for a in args if isinstance(a, str)), None)
                if q:
                    span.set_attribute("input.value", str(q))
                result = fn(*args, **kwargs)
                if isinstance(result, list) and result and isinstance(result[0], Document):
                    span.set_attribute("retrieval.documents", encode_documents(result))
                    span.set_attribute("output.value", _document_summary(result))
                span.set_status(Status(StatusCode.OK))
                return result

        return wrapper  # type: ignore[return-value]

    return decorator


class RetrievalSpan:
    """Context manager for instrumenting retrieval logic as a RETRIEVER span.

    Use when you need fine-grained control over when documents and query
    are attached (e.g. streaming or multi-step retrieval).

    Example::

        with RetrievalSpan("my_retrieval") as r:
            r.set_query(query)
            docs = fetch_documents(query)
            r.set_documents(docs)
    """

    def __init__(self, name: str = "retrieval") -> None:
        self._name = name
        self._span: Any = None
        self._ctx: Any = None

    def __enter__(self) -> RetrievalSpan:
        tracer = _get_tracer()
        self._ctx = tracer.start_as_current_span(self._name)
        self._span = self._ctx.__enter__()
        self._span.set_attribute("openinference.span.kind", "RETRIEVER")
        return self

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

    def set_documents(self, docs: list[Document]) -> None:
        self._span.set_attribute("retrieval.documents", encode_documents(docs))
        self._span.set_attribute("output.value", _document_summary(docs))

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool | None:
        if exc_val is None:
            self._span.set_status(Status(StatusCode.OK))
        return self._ctx.__exit__(exc_type, exc_val, exc_tb)
