# eval_lib/tracing/langchain_callback.py
"""LangChain callback handler for automatic tracing"""

from typing import Any, Dict, List, Optional, Union
from uuid import UUID

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.messages import BaseMessage
from langchain_core.outputs import LLMResult

from .tracer import tracer
from .types import SpanType


class EvalLibCallbackHandler(BaseCallbackHandler):
    """LangChain callback handler that creates spans for all LangChain operations"""

    def __init__(self):
        super().__init__()
        # Map run_id -> span for tracking nested spans
        self._run_spans: Dict[UUID, Any] = {}

    # ============ Chain callbacks ============

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a chain starts running"""
        if serialized is None:
            name = "Chain"
        else:
            name = serialized.get("name") or serialized.get("id", ["Unknown"])[-1]

        span = tracer.start_span(
            name=name,
            span_type=SpanType.AGENT_STEP,
            input_data=inputs,
            metadata={"tags": tags, **(metadata or {})}
        )
        if span:
            self._run_spans[run_id] = span

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Called when a chain finishes"""
        span = self._run_spans.pop(run_id, None)
        if span:
            tracer.end_span(span, output=outputs)

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Called when a chain errors"""
        span = self._run_spans.pop(run_id, None)
        if span:
            tracer.end_span(span, error=error)

    # ============ LLM callbacks ============

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when LLM starts"""
        if serialized is None:
            name = "LLM"
        else:
            name = serialized.get("name") or serialized.get("id", ["LLM"])[-1]

        span = tracer.start_span(
            name=name,
            span_type=SpanType.LLM_CALL,
            input_data={"prompts": prompts},
            metadata={"tags": tags, **(metadata or {})}
        )
        if span:
            self._run_spans[run_id] = span

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when chat model starts"""
        if serialized is None:
            name = "ChatModel"
        else:
            name = serialized.get("name") or serialized.get("id", ["ChatModel"])[-1]

        # Convert messages to serializable format
        messages_data = []
        for msg_list in messages:
            messages_data.append([
                {"type": msg.type, "content": msg.content}
                for msg in msg_list
            ])

        span = tracer.start_span(
            name=name,
            span_type=SpanType.LLM_CALL,
            input_data={"messages": messages_data},
            metadata={"tags": tags, **(metadata or {})}
        )
        if span:
            self._run_spans[run_id] = span

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Called when LLM ends"""
        span = self._run_spans.pop(run_id, None)
        if span:
            # Extract output text
            output_data = {
                "generations": [
                    [{"text": gen.text, "message": getattr(gen, "message", None)}
                     for gen in gens]
                    for gens in response.generations
                ]
            }
            if response.llm_output:
                output_data["llm_output"] = response.llm_output

            tracer.end_span(span, output=output_data)

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Called when LLM errors"""
        span = self._run_spans.pop(run_id, None)
        if span:
            tracer.end_span(span, error=error)

    # ============ Tool callbacks ============

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        inputs: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a tool starts"""
        if serialized is None:
            name = "Tool"
        else:
            name = serialized.get("name", "Tool")

        span = tracer.start_span(
            name=name,
            span_type=SpanType.TOOL_CALL,
            input_data=inputs or {"input": input_str},
            metadata={"tags": tags, **(metadata or {})}
        )
        if span:
            self._run_spans[run_id] = span

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Called when a tool ends"""
        span = self._run_spans.pop(run_id, None)
        if span:
            tracer.end_span(span, output=output)

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Called when a tool errors"""
        span = self._run_spans.pop(run_id, None)
        if span:
            tracer.end_span(span, error=error)

    # ============ Agent callbacks ============

    def on_agent_action(
        self,
        action: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when agent takes an action"""
        # Agent actions are typically followed by tool calls
        # which we track separately
        pass

    def on_agent_finish(
        self,
        finish: Any,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Called when agent finishes"""
        pass

    # ============ Retriever callbacks ============

    def on_retriever_start(
        self,
        serialized: Dict[str, Any],
        query: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when retriever starts"""
        if serialized is None:
            name = "Retriever"
        else:
            name = serialized.get("name", "Retriever")

        span = tracer.start_span(
            name=name,
            span_type=SpanType.RETRIEVAL,
            input_data={"query": query},
            metadata={"tags": tags, **(metadata or {})}
        )
        if span:
            self._run_spans[run_id] = span

    def on_retriever_end(
        self,
        documents: List[Any],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Called when retriever ends"""
        span = self._run_spans.pop(run_id, None)
        if span:
            tracer.end_span(span, output={"documents": [doc.page_content for doc in documents]})

    def on_retriever_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        """Called when retriever errors"""
        span = self._run_spans.pop(run_id, None)
        if span:
            tracer.end_span(span, error=error)


# Singleton instance for easy import
callback_handler = EvalLibCallbackHandler()