"""Reranker span decorator and context manager."""

from __future__ import annotations

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

from coalex.ext._base import _get_tracer
from coalex.ext.retrieval import Document, encode_documents

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


def reranker_span(
    name: str | None = None,
    model_name: str | None = None,
    query_arg: str = "query",
    docs_arg: str = "documents",
    top_k: int | None = None,
) -> Callable[[F], F]:
    """Decorator that wraps a reranker function with an OpenInference RERANKER span.

    Supports both sync and async functions. The decorated function must accept a
    ``list[Document]`` as input and return ``list[Document]`` as output. The span captures:
    - ``openinference.span.kind`` = ``RERANKER``
    - ``reranker.model_name`` — if provided
    - ``reranker.query`` — from ``query_arg`` kwarg or first ``str`` positional
    - ``reranker.top_k`` — if provided
    - ``reranker.input_documents`` — JSON-encoded input docs
    - ``reranker.output_documents`` — JSON-encoded output docs

    Args:
        name: Span name. Defaults to the function name.
        model_name: Reranker model name to record.
        query_arg: Name of the kwarg containing the query string.
        docs_arg: Name of the kwarg containing the input documents list.
        top_k: Number of top documents to return (recorded as attribute).
    """

    def _set_span_attrs(span: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None:
        span.set_attribute("openinference.span.kind", "RERANKER")
        if model_name:
            span.set_attribute("reranker.model_name", model_name)
        if top_k is not None:
            span.set_attribute("reranker.top_k", top_k)
        q = kwargs.get(query_arg) or next((a for a in args if isinstance(a, str)), None)
        if q:
            span.set_attribute("reranker.query", str(q))
        input_docs = kwargs.get(docs_arg) or next((a for a in args if isinstance(a, list)), None)
        if input_docs and input_docs and isinstance(input_docs[0], Document):
            span.set_attribute("reranker.input_documents", encode_documents(input_docs))

    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:
                    _set_span_attrs(span, args, kwargs)
                    result = await fn(*args, **kwargs)
                    if isinstance(result, list) and result and isinstance(result[0], Document):
                        span.set_attribute("reranker.output_documents", encode_documents(result))
                    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:
                _set_span_attrs(span, args, kwargs)
                result = fn(*args, **kwargs)
                if isinstance(result, list) and result and isinstance(result[0], Document):
                    span.set_attribute("reranker.output_documents", encode_documents(result))
                return result

        return wrapper  # type: ignore[return-value]

    return decorator


class RerankerSpan:
    """Context manager for instrumenting a reranking step as a RERANKER span.

    Example::

        with RerankerSpan("rerank", model_name="rerank-v3", top_k=5) as r:
            r.set_query(query)
            r.set_input_documents(docs)
            ranked = reranker.rerank(query, docs, top_k=5)
            r.set_output_documents(ranked)
    """

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

    def __enter__(self) -> RerankerSpan:
        tracer = _get_tracer()
        self._ctx = tracer.start_as_current_span(self._name)
        self._span = self._ctx.__enter__()
        self._span.set_attribute("openinference.span.kind", "RERANKER")
        if self._model_name:
            self._span.set_attribute("reranker.model_name", self._model_name)
        if self._top_k is not None:
            self._span.set_attribute("reranker.top_k", self._top_k)
        return self

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

    def set_input_documents(self, docs: list[Document]) -> None:
        self._span.set_attribute("reranker.input_documents", encode_documents(docs))

    def set_output_documents(self, docs: list[Document]) -> None:
        self._span.set_attribute("reranker.output_documents", encode_documents(docs))

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