"""
LangChain Callback Handler for Pendo SDK

Creates OTel spans from LangChain callbacks. Every LLM call, tool call,
and chain execution becomes a traced span that flows to Pendo.

Usage:
    from pendo_sdk.langchain_callback import PendoCallbackHandler
    handler = PendoCallbackHandler()

    # Add to LangChain as a global callback:
    import langchain
    langchain.callbacks.manager.add_handler(handler)

    # Or pass per-call:
    llm.invoke(prompt, config={"callbacks": [handler]})
"""

import logging
import time
import uuid
from typing import Any, Dict, List, Optional, Union

from opentelemetry import trace, context
from opentelemetry.trace import StatusCode

logger = logging.getLogger("pendo_sdk")

_SDK_NAME = "pendo_sdk"


class PendoCallbackHandler:
    """LangChain callback handler that creates OTel spans for Pendo tracing."""

    # Required attributes for LangChain callback manager compatibility
    run_inline = True
    raise_error = False
    ignore_chain = False
    ignore_llm = False
    ignore_retry = True
    ignore_agent = False
    ignore_chat_model = False
    ignore_retriever = True
    ignore_custom_event = True

    # Chain names to suppress — these are LangGraph infrastructure, not meaningful traces
    _NOISE_CHAINS = {"start_node", "route_after_start_node", "__start__", "next_actions_node",
                     "route_after_supervisor", "RunnableSequence", "RunnableLambda", ""}

    def __init__(self, conversation_id: Optional[str] = None):
        self._tracer = trace.get_tracer(_SDK_NAME)
        self._spans: Dict[str, Any] = {}  # run_id -> (span, token)
        self._conversation_id = conversation_id or ""
        self._root_span = None  # Root span for the current conversation turn
        self._root_token = None

    def set_conversation_id(self, conversation_id: str):
        self._conversation_id = conversation_id

    def start_turn(self):
        """Start a new conversation turn — creates a root span that all child spans nest under.
        This ensures all spans in a single turn share the same agentTraceId.
        """
        self.end_turn()  # Clean up any previous turn
        self._root_span = self._tracer.start_span("conversation.turn")
        self._root_span.set_attribute("pendo.span.type", "AGENT")
        if self._conversation_id:
            self._root_span.set_attribute("pendo.conversation_id", self._conversation_id)
        self._root_token = context.attach(trace.set_span_in_context(self._root_span))

    def end_turn(self):
        """End the current conversation turn."""
        if self._root_span:
            self._root_span.end()
            if self._root_token:
                try:
                    context.detach(self._root_token)
                except ValueError:
                    pass
            self._root_span = None
            self._root_token = None

    def _start_span(self, run_id: str, name: str, span_type: str,
                    attrs: Optional[Dict] = None, parent_run_id: Optional[str] = None):
        # Determine parent context:
        # 1. If parent_run_id maps to an active span, use that
        # 2. Otherwise use the root turn span (so all spans share same trace ID)
        # 3. If no root span, create one automatically
        parent_ctx = None
        if parent_run_id:
            parent_key = str(parent_run_id)
            if parent_key in self._spans:
                parent_span, _ = self._spans[parent_key]
                parent_ctx = trace.set_span_in_context(parent_span)

        if not parent_ctx and self._root_span:
            parent_ctx = trace.set_span_in_context(self._root_span)

        if not parent_ctx:
            # Auto-create root span if none exists
            self.start_turn()
            parent_ctx = trace.set_span_in_context(self._root_span)

        span = self._tracer.start_span(name, context=parent_ctx)
        span.set_attribute("pendo.span.type", span_type)
        if self._conversation_id:
            span.set_attribute("pendo.conversation_id", self._conversation_id)
        if attrs:
            for k, v in attrs.items():
                if v is not None:
                    span.set_attribute(k, str(v) if not isinstance(v, (int, float, bool)) else v)
        token = context.attach(trace.set_span_in_context(span))
        self._spans[str(run_id)] = (span, token)

    def _emit_prompt_event(self, content: str, model: str = ""):
        """Emit a standalone prompt span for the user's message.

        This creates a separate short-lived span so the user's prompt
        appears as a 'prompt' conversation event, independent of the
        LLM generation span which will become 'agent_response'.
        """
        parent_ctx = trace.set_span_in_context(self._root_span) if self._root_span else None
        span = self._tracer.start_span("user.prompt", context=parent_ctx)
        span.set_attribute("pendo.span.type", "USER_MESSAGE")
        span.set_attribute("message.role", "user")
        span.set_attribute("message.content", content[:500] if content else "")
        if model:
            span.set_attribute("llm.model_name", model)
        if self._conversation_id:
            span.set_attribute("pendo.conversation_id", self._conversation_id)
        span.end()  # End immediately — this is a point event

    def _end_span(self, run_id: str, error: Optional[str] = None):
        key = str(run_id)
        if key not in self._spans:
            return
        span, token = self._spans.pop(key)
        if error:
            span.set_status(StatusCode.ERROR, error)
        else:
            span.set_status(StatusCode.OK)
        span.end()
        try:
            context.detach(token)
        except ValueError:
            pass

    # ---- LLM callbacks ----

    def _format_messages_as_prompt(self, messages: List[Any]) -> str:
        """Format LangChain messages into a complete prompt string for agentTraceContent.

        Captures the full message chain including system prompt, conversation
        history, and user message — this is what the LLM actually received.
        """
        if not messages:
            return ""
        msg_list = messages[0] if isinstance(messages[0], list) else messages
        parts = []
        for msg in msg_list:
            role = getattr(msg, "type", "unknown")
            content = str(getattr(msg, "content", ""))
            if not content:
                continue
            parts.append(f"[{role}] {content}")
        return "\n".join(parts)

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        serialized = serialized or {}
        serialized_kwargs = (serialized.get("kwargs") or {})
        model = (
            serialized_kwargs.get("model_name", "")
            or serialized_kwargs.get("model", "")
            or serialized.get("id", [""])[-1]
        )
        content = "\n".join(p[:300] for p in prompts)[:500] if prompts else ""
        self._start_span(run_id, f"llm.{model}", "GENERATION", {
            "llm.model_name": model,
            "message.content": content,
            "message.role": "user",
        }, parent_run_id=parent_run_id)

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[Any],
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        serialized = serialized or {}
        serialized_kwargs = serialized.get("kwargs") or {}
        model = (
            serialized_kwargs.get("model_name", "")
            or serialized_kwargs.get("model", "")
            or kwargs.get("invocation_params", {}).get("model_name", "")
            or kwargs.get("invocation_params", {}).get("model", "")
            or serialized.get("id", [""])[-1]
        )
        # Full prompt context for agentTraceContent
        full_prompt = self._format_messages_as_prompt(messages)
        # User's actual message for conversation event content
        user_content = ""
        msg_list = messages[0] if messages and isinstance(messages[0], list) else (messages or [])
        for msg in reversed(msg_list):
            if getattr(msg, "type", "") == "human":
                user_content = str(getattr(msg, "content", ""))[:500]
                break
        # Emit a separate prompt event for the user's message
        if user_content:
            self._emit_prompt_event(user_content, model)

        # The LLM generation span starts as "assistant" — it'll become agent_response
        # when on_llm_end fires with the model's response content
        self._start_span(run_id, f"llm.{model}", "GENERATION", {
            "llm.model_name": model,
            "message.content": "",  # Will be set by on_llm_end with response content
            "message.trace_content": full_prompt,
            "message.role": "assistant",
        }, parent_run_id=parent_run_id)

    def on_llm_end(
        self,
        response: Any,
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        **kwargs: Any,
    ) -> None:
        key = str(run_id)
        if key in self._spans:
            span, _ = self._spans[key]
            # Extract token usage
            llm_output = getattr(response, "llm_output", None) or {}
            usage = llm_output.get("token_usage", {})
            if usage:
                if "prompt_tokens" in usage:
                    span.set_attribute("llm.token_count.prompt", usage["prompt_tokens"])
                if "completion_tokens" in usage:
                    span.set_attribute("llm.token_count.completion", usage["completion_tokens"])
                if "total_tokens" in usage:
                    span.set_attribute("llm.token_count.total", usage["total_tokens"])
            model = llm_output.get("model_name", "")
            if model:
                span.set_attribute("llm.model_name", model)

            # Extract the LLM's response text for agent_response event
            response_text = ""
            generations = getattr(response, "generations", None)
            if generations and generations[0]:
                gen = generations[0][0] if isinstance(generations[0], list) else generations[0]
                response_text = str(getattr(gen, "text", ""))[:500]
                # Check for tool calls in the response
                msg = getattr(gen, "message", None)
                if msg:
                    tool_calls = getattr(msg, "tool_calls", None)
                    if tool_calls:
                        tool_names = [tc.get("name", "") for tc in tool_calls if isinstance(tc, dict)]
                        if tool_names:
                            span.set_attribute("tool.name", tool_names[0])
                            response_text = response_text or f"tool_call: {', '.join(tool_names)}"

            # Set response content — exporter uses this for agent_response event content
            if response_text:
                span.set_attribute("message.content", response_text)
                span.set_attribute("message.response_content", response_text)

        self._end_span(run_id)

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        **kwargs: Any,
    ) -> None:
        self._end_span(run_id, error=str(error))

    # ---- Tool callbacks ----

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        tool_name = (serialized or {}).get("name", "") or (serialized or {}).get("id", [""])[-1]
        # Detect subagent calls (tools that delegate to sub-agents)
        is_subagent = tool_name.startswith("call_") or tool_name.startswith("handoff_")
        span_type = "CHAIN" if is_subagent else "TOOL"
        trace_type_override = "subagent_request" if is_subagent else None

        self._start_span(run_id, f"tool.{tool_name}", span_type, {
            "tool.name": tool_name,
            "tool.input": str(input_str)[:500],
        }, parent_run_id=parent_run_id)

        # Override trace type for subagent calls
        if trace_type_override:
            key = str(run_id)
            if key in self._spans:
                span, _ = self._spans[key]
                span.set_attribute("pendo.trace_type_override", trace_type_override)

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        **kwargs: Any,
    ) -> None:
        key = str(run_id)
        if key in self._spans:
            span, _ = self._spans[key]
            span.set_attribute("tool.output", str(output)[:500])
        self._end_span(run_id)

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        **kwargs: Any,
    ) -> None:
        self._end_span(run_id, error=str(error))

    # ---- Chain callbacks ----

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        chain_name = (
            (metadata or {}).get("langgraph_node", "")
            or (serialized or {}).get("name", "")
            or (serialized or {}).get("id", [""])[-1]
        )

        # Skip noise chains — LangGraph infrastructure spans
        if chain_name in self._NOISE_CHAINS:
            return

        self._start_span(run_id, f"chain.{chain_name}", "CHAIN",
                         parent_run_id=parent_run_id)

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        **kwargs: Any,
    ) -> None:
        self._end_span(run_id)

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: Any,
        parent_run_id: Optional[Any] = None,
        **kwargs: Any,
    ) -> None:
        self._end_span(run_id, error=str(error))

    # ---- Required no-op callbacks ----

    def on_text(self, text: str, **kwargs: Any) -> None:
        pass

    def on_agent_action(self, action: Any, **kwargs: Any) -> None:
        pass

    def on_agent_finish(self, finish: Any, **kwargs: Any) -> None:
        pass

    def on_retry(self, retry_state: Any, **kwargs: Any) -> None:
        pass

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        pass
