"""
Retrieval tracking for RAG systems.

Provides decorators and context managers to track retrieval operations
in Langgraph tools and other custom retrieval implementations.

Uses OpenInference semantic conventions for RETRIEVER spans:
- openinference.span.kind = "RETRIEVER"
- retrieval.query = the search query
- retrieval.documents = list of retrieved documents
- retrieval.top_k = number of documents

Example usage:

    # Decorator pattern
    @coalex.retrieval()
    def search_documents(query: str) -> List[Document]:
        return vector_db.search(query)

    # Context manager pattern (more control)
    with coalex.retrieval_context(query="What is RAG?") as ctx:
        docs = vector_db.search(query)
        ctx.set_documents(docs)

    # Async decorator
    @coalex.async_retrieval()
    async def async_search(query: str) -> List[Document]:
        return await vector_db.asearch(query)
"""

import functools
import inspect
import json
import logging
import threading
from contextlib import contextmanager
from dataclasses import dataclass, field, asdict
from typing import Any, Callable, Dict, List, Optional, TypeVar, Union

from opentelemetry import trace
from opentelemetry.trace import SpanKind as OTelSpanKind

logger = logging.getLogger(__name__)

# OpenInference span kind for retrievers
OPENINFERENCE_SPAN_KIND_RETRIEVER = "RETRIEVER"

# Thread-local storage for current retrieval context
_retrieval_context_var = threading.local()


@dataclass
class Document:
    """
    Represents a retrieved document.

    Follows OpenInference semantic conventions for document representation.

    Attributes:
        content: The document text content (required)
        id: Optional document identifier
        score: Optional relevance/similarity score
        metadata: Optional additional metadata
    """

    content: str
    id: Optional[str] = None
    score: Optional[float] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_openinference_dict(self) -> Dict[str, Any]:
        """Convert to OpenInference document format."""
        doc = {"document.content": self.content}
        if self.id:
            doc["document.id"] = self.id
        if self.score is not None:
            doc["document.score"] = self.score
        if self.metadata:
            doc["document.metadata"] = self.metadata
        return doc


def _normalize_document(doc: Any) -> Dict[str, Any]:
    """
    Normalize various document formats to OpenInference format.

    Supports:
    - Document dataclass
    - Dict with various key patterns
    - LangChain Document objects (page_content)
    - LlamaIndex NodeWithScore objects
    - Plain strings
    """
    if isinstance(doc, Document):
        return doc.to_openinference_dict()

    if isinstance(doc, dict):
        # Already in OpenInference format
        if "document.content" in doc:
            return doc

        # Common dict patterns
        content = (
            doc.get("content")
            or doc.get("page_content")  # LangChain
            or doc.get("text")
            or doc.get("body")
            or str(doc)
        )

        result = {"document.content": content}

        # Extract ID
        doc_id = doc.get("id") or doc.get("doc_id") or doc.get("document_id")
        if doc_id:
            result["document.id"] = str(doc_id)

        # Extract score
        score = doc.get("score") or doc.get("similarity") or doc.get("relevance_score")
        if score is not None:
            try:
                result["document.score"] = float(score)
            except (ValueError, TypeError):
                pass

        # Extract metadata
        metadata = doc.get("metadata") or doc.get("meta")
        if metadata and isinstance(metadata, dict):
            result["document.metadata"] = metadata

        return result

    # LangChain Document object
    if hasattr(doc, "page_content"):
        result = {"document.content": doc.page_content}
        if hasattr(doc, "metadata") and doc.metadata:
            result["document.metadata"] = doc.metadata
        return result

    # LlamaIndex NodeWithScore
    if hasattr(doc, "node") and hasattr(doc, "score"):
        node = doc.node
        result = {
            "document.content": getattr(node, "text", str(node)),
            "document.score": float(doc.score) if doc.score else None,
        }
        if hasattr(node, "node_id"):
            result["document.id"] = node.node_id
        if hasattr(node, "metadata"):
            result["document.metadata"] = node.metadata
        return result

    # Fallback: convert to string
    return {"document.content": str(doc)}


class RetrievalContext:
    """
    Context for tracking retrieval operations.

    Provides methods to set query and documents on the current span.
    Used with the retrieval_context() context manager.
    """

    def __init__(self, span: trace.Span, query: Optional[str] = None):
        """
        Initialize retrieval context.

        Args:
            span: The OpenTelemetry span to add attributes to
            query: Optional initial query
        """
        self._span = span
        self._query = query
        self._documents: List[Dict[str, Any]] = []

        if query:
            self._span.set_attribute("retrieval.query", query)

    @property
    def span(self) -> trace.Span:
        """Get the underlying span."""
        return self._span

    def set_query(self, query: str) -> "RetrievalContext":
        """
        Set the retrieval query.

        Args:
            query: The search query string

        Returns:
            Self for chaining
        """
        self._query = query
        self._span.set_attribute("retrieval.query", query)
        return self

    def set_documents(
        self, documents: Union[List[Document], List[Dict[str, Any]], List[Any]]
    ) -> "RetrievalContext":
        """
        Set the retrieved documents.

        Args:
            documents: List of documents in various formats

        Returns:
            Self for chaining
        """
        self._documents = [_normalize_document(doc) for doc in documents]

        # Serialize documents to JSON for span attribute
        try:
            docs_json = json.dumps(self._documents)
            self._span.set_attribute("retrieval.documents", docs_json)
        except (TypeError, ValueError) as e:
            logger.warning(f"Failed to serialize documents: {e}")
            # Fallback: store document count only
            self._span.set_attribute("retrieval.documents", "[]")

        self._span.set_attribute("retrieval.top_k", len(self._documents))

        # Extract and set scores if available
        scores = [
            doc.get("document.score")
            for doc in self._documents
            if doc.get("document.score") is not None
        ]
        if scores:
            self._span.set_attribute("retrieval.scores", json.dumps(scores))

        return self

    def add_document(
        self,
        content: str,
        score: Optional[float] = None,
        doc_id: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> "RetrievalContext":
        """
        Add a single document to the retrieval results.

        Args:
            content: Document text content
            score: Optional relevance score
            doc_id: Optional document identifier
            metadata: Optional metadata dict

        Returns:
            Self for chaining
        """
        doc = Document(
            content=content, id=doc_id, score=score, metadata=metadata or {}
        )
        self._documents.append(doc.to_openinference_dict())

        # Update span attributes
        try:
            docs_json = json.dumps(self._documents)
            self._span.set_attribute("retrieval.documents", docs_json)
        except (TypeError, ValueError):
            pass

        self._span.set_attribute("retrieval.top_k", len(self._documents))
        return self


def _get_current_retrieval_context() -> Optional["RetrievalContext"]:
    """Get the current retrieval context from thread-local storage."""
    return getattr(_retrieval_context_var, "context", None)


def _set_current_retrieval_context(ctx: Optional["RetrievalContext"]) -> None:
    """Set the current retrieval context in thread-local storage."""
    _retrieval_context_var.context = ctx


def track_documents(
    documents: Union[List[Document], List[Dict[str, Any]], List[Any]]
) -> None:
    """
    Track retrieved documents from within a @retrieval() decorated function.

    This allows you to explicitly set the documents when your function
    returns something other than a list (e.g., a formatted string).

    Args:
        documents: List of documents in various formats

    Example:
        @tool
        @coalex.retrieval()
        def search_knowledge_base(query: str) -> str:
            docs = vector_db.search(query)
            coalex.track_documents(docs)  # Track the actual documents
            return format_docs_for_llm(docs)  # Return formatted string
    """
    ctx = _get_current_retrieval_context()
    if ctx is None:
        logger.warning(
            "track_documents() called outside of a retrieval context. "
            "Make sure to use @coalex.retrieval() decorator or retrieval_context()."
        )
        return

    ctx.set_documents(documents)


@contextmanager
def retrieval_context(
    query: Optional[str] = None,
    span_name: str = "retrieval",
    top_k: Optional[int] = None,
):
    """
    Context manager for tracking retrieval operations.

    Creates a RETRIEVER span with OpenInference attributes.
    Use this for fine-grained control over retrieval tracking.

    Args:
        query: The search query (can also be set later via context.set_query())
        span_name: Name for the span (default: "retrieval")
        top_k: Expected number of documents to retrieve (optional)

    Yields:
        RetrievalContext for setting documents and additional attributes

    Example:
        with coalex.retrieval_context(query="What is machine learning?") as ctx:
            docs = vector_db.search(query)
            ctx.set_documents(docs)

        # Or set query later
        with coalex.retrieval_context() as ctx:
            query = process_user_input(input)
            ctx.set_query(query)
            docs = search(query)
            ctx.set_documents(docs)
    """
    tracer = trace.get_tracer(__name__)

    with tracer.start_as_current_span(span_name, kind=OTelSpanKind.INTERNAL) as span:
        # Set OpenInference span kind
        span.set_attribute("openinference.span.kind", OPENINFERENCE_SPAN_KIND_RETRIEVER)

        if top_k is not None:
            span.set_attribute("retrieval.top_k", top_k)

        ctx = RetrievalContext(span, query)

        # Store in thread-local for track_documents() access
        previous_ctx = _get_current_retrieval_context()
        _set_current_retrieval_context(ctx)

        try:
            yield ctx
        finally:
            # Restore previous context
            _set_current_retrieval_context(previous_ctx)


# Type variable for preserving function signatures
F = TypeVar("F", bound=Callable[..., Any])


def retrieval(
    query_param: str = "query",
    documents_from_return: bool = True,
    span_name: Optional[str] = None,
) -> Callable[[F], F]:
    """
    Decorator for tracking retrieval operations.

    Automatically creates a RETRIEVER span, extracts the query from
    function parameters, and captures returned documents.

    Args:
        query_param: Name of the parameter containing the query (default: "query")
        documents_from_return: If True, treat return value as documents (default: True)
        span_name: Optional custom span name (defaults to function name)

    Returns:
        Decorated function

    Example:
        @coalex.retrieval()
        def search_documents(query: str) -> List[Document]:
            docs = vector_db.search(query)
            return docs

        @coalex.retrieval(query_param="search_term")
        def custom_search(search_term: str, filters: dict) -> List[dict]:
            return db.search(search_term, filters)

        # In a Langgraph tool
        @tool
        @coalex.retrieval()
        def rag_search(query: str) -> str:
            docs = search_documents(query)
            return format_docs(docs)
    """

    def decorator(func: F) -> F:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # Determine span name
            name = span_name or func.__name__

            # Extract query from parameters
            query = _extract_query_from_args(func, args, kwargs, query_param)

            with retrieval_context(query=query, span_name=name) as ctx:
                result = func(*args, **kwargs)

                if documents_from_return and result is not None:
                    _set_documents_from_result(ctx, result)

                return result

        return wrapper

    return decorator


def async_retrieval(
    query_param: str = "query",
    documents_from_return: bool = True,
    span_name: Optional[str] = None,
) -> Callable[[F], F]:
    """
    Async decorator for tracking retrieval operations.

    Same as @retrieval() but for async functions.

    Args:
        query_param: Name of the parameter containing the query (default: "query")
        documents_from_return: If True, treat return value as documents (default: True)
        span_name: Optional custom span name (defaults to function name)

    Returns:
        Decorated async function

    Example:
        @coalex.async_retrieval()
        async def async_search(query: str) -> List[Document]:
            return await vector_db.asearch(query)
    """

    def decorator(func: F) -> F:
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            name = span_name or func.__name__
            query = _extract_query_from_args(func, args, kwargs, query_param)

            with retrieval_context(query=query, span_name=name) as ctx:
                result = await func(*args, **kwargs)

                if documents_from_return and result is not None:
                    _set_documents_from_result(ctx, result)

                return result

        return wrapper

    return decorator


def _extract_query_from_args(
    func: Callable, args: tuple, kwargs: dict, query_param: str
) -> Optional[str]:
    """Extract query value from function arguments."""
    # Check kwargs first
    if query_param in kwargs:
        return str(kwargs[query_param])

    # Try to get from positional args using function signature
    try:
        sig = inspect.signature(func)
        params = list(sig.parameters.keys())
        if query_param in params:
            idx = params.index(query_param)
            if idx < len(args):
                return str(args[idx])
    except (ValueError, TypeError):
        pass

    return None


def _set_documents_from_result(ctx: RetrievalContext, result: Any) -> None:
    """Set documents on context from function return value."""
    if isinstance(result, list):
        ctx.set_documents(result)
    elif hasattr(result, "__iter__") and not isinstance(result, (str, bytes)):
        try:
            ctx.set_documents(list(result))
        except TypeError:
            pass
