"""
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__", "__end__",
        "next_actions_node", "route_after_supervisor", "route_after_start",
        "RunnableSequence", "RunnableLambda", "RunnableParallel",
        "ChannelWrite", "ChannelRead", "",
    }

    def __init__(self, conversation_id: Optional[str] = None):
        self._tracer = trace.get_tracer(_SDK_NAME)
        self._prompt_emitted = False  # Only emit one prompt event per turn
        self._spans: Dict[str, Any] = {}  # run_id -> (span, token)
        self._chain_parents: Dict[str, Optional[str]] = {}  # chain run_id -> parent run_id
        self._conversation_id = conversation_id or ""
        self._root_span = None  # Root span for the current conversation turn
        self._root_token = None
        # Accumulate response data across the turn for a single merged agent_response
        self._turn_response_text = ""  # Last response text wins
        self._turn_models: List[str] = []
        self._turn_tools: List[str] = []

    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._prompt_emitted = False  # Reset prompt flag for new turn
        self._turn_response_text = ""
        self._turn_models = []
        self._turn_tools = []
        self._root_span = self._tracer.start_span("conversation.turn")
        self._root_span.set_attribute("pendo.span.type", "AGENT")
        self._root_span.set_attribute("message.role", "system")  # Won't emit conversation event
        if self._conversation_id:
            self._root_span.set_attribute("pendo.conversation_id", self._conversation_id)

    def end_turn(self):
        """End the current conversation turn.
        Emits a single merged agent_response conversation event with accumulated
        models and tools from all LLM calls in the turn.
        """
        if self._root_span:
            # Emit one merged agent_response if we have response text
            if self._turn_response_text:
                parent_ctx = trace.set_span_in_context(self._root_span)
                span = self._tracer.start_span("agent.response", context=parent_ctx)
                span.set_attribute("pendo.span.type", "AGENT_RESPONSE")
                span.set_attribute("message.role", "assistant")
                span.set_attribute("message.content", self._turn_response_text)
                span.set_attribute("message.response_content", self._turn_response_text)
                if self._turn_models:
                    span.set_attribute("llm.model_name", ", ".join(self._turn_models))
                if self._turn_tools:
                    span.set_attribute("tool.name", ", ".join(self._turn_tools))
                if self._conversation_id:
                    span.set_attribute("pendo.conversation_id", self._conversation_id)
                span.end()

            self._root_span.end()
            self._root_span = None

    def _resolve_parent_run_id(self, parent_run_id: Optional[str]) -> Optional[str]:
        """Traverse chain parents to find an ancestor that has a real span.
        Chains don't emit events but sub-agent LLM/tool calls are parented to them.
        Walk up the chain_parents map until we find a run_id in self._spans.
        """
        if not parent_run_id:
            return None
        current = str(parent_run_id)
        visited = set()
        while current and current not in visited:
            if current in self._spans:
                return current
            visited.add(current)
            current = self._chain_parents.get(current)
        return 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. Resolve parent through chain hierarchy to find a real span
        # 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
        resolved_parent = self._resolve_parent_run_id(str(parent_run_id) if parent_run_id else None)
        if resolved_parent and resolved_parent in self._spans:
            parent_span, _ = self._spans[resolved_parent]
            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)
        # Don't attach to global context — async tasks cross boundaries
        # and context.detach fails across different async contexts.
        # Just store the span for parent lookup.
        self._spans[str(run_id)] = (span, None)

    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 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, _ = self._spans.pop(key)
        if error:
            span.set_status(StatusCode.ERROR, error)
        else:
            span.set_status(StatusCode.OK)
        span.end()

    # ---- LLM callbacks ----

    def _format_messages_as_prompt(self, messages: List[Any]) -> str:
        """Format the delta messages for this specific LLM call's agentTraceContent.

        Only captures NEW messages since the last assistant response — not the
        full conversation history. System prompts, previous tool calls, and
        previous AI responses already have their own trace events.

        For the first LLM call in a turn: includes system + user message.
        For subsequent calls (after tool use): includes only tool results + new user messages.
        """
        if not messages:
            return ""
        msg_list = messages[0] if isinstance(messages[0], list) else messages

        # Find the last assistant message index — everything after it is "new"
        last_ai_idx = -1
        for i, msg in enumerate(msg_list):
            role = getattr(msg, "type", "unknown")
            if role in ("ai", "assistant"):
                last_ai_idx = i

        # If no prior AI message, this is the first LLM call — include system + user
        if last_ai_idx == -1:
            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)

        # Otherwise, only include messages AFTER the last AI response (the delta)
        parts = []
        for msg in msg_list[last_ai_idx + 1:]:
            role = getattr(msg, "type", "unknown")
            content = str(getattr(msg, "content", ""))
            if not content:
                continue
            parts.append(f"[{role}] {content}")
        return "\n".join(parts) if parts else ""

    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(prompts) 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", ""))
                break
        # Emit a separate prompt event for the user's message
        # Always emit — use user_content if found, otherwise extract from full prompt
        prompt_text = user_content
        if not prompt_text and full_prompt:
            # Try to extract last [human] message from formatted prompt
            for line in reversed(full_prompt.split("\n")):
                if line.startswith("[human]"):
                    prompt_text = line[len("[human]"):].strip()
                    break
        if prompt_text and not self._prompt_emitted:
            self._emit_prompt_event(prompt_text, model)
            self._prompt_emitted = True
            logger.info("Emitted prompt event: %s", prompt_text[:80])
        else:
            logger.warning("No user content found for prompt event. msg_list types: %s",
                          [type(m).__name__ for m in msg_list[:3]])

        # 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", ""))
                # 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)}"

            # Determine if this is a tool call or a real text response
            has_tool_calls = bool(getattr(getattr(gen, "message", None), "tool_calls", None) if generations and generations[0] else None)

            if has_tool_calls:
                # Tool call decision — LLM decided to call a tool/agent.
                # Capture the full response including reasoning + tool call info.
                span.set_attribute("message.role", "tool_decision")

                # Build the decision output: reasoning text + tool calls
                decision_parts = []
                # Get reasoning from message.content (where LLMs put their thinking)
                if msg:
                    msg_content = str(getattr(msg, "content", ""))
                    if msg_content:
                        decision_parts.append(msg_content)
                # Append tool call details
                tool_calls_list = getattr(msg, "tool_calls", []) if msg else []
                for tc in tool_calls_list:
                    if isinstance(tc, dict):
                        tc_name = tc.get("name", "")
                        tc_args = tc.get("args", {})
                        decision_parts.append(f"Action: {tc_name}({tc_args})")
                        if tc_name and tc_name not in self._turn_tools:
                            self._turn_tools.append(tc_name)

                decision_text = "\n".join(decision_parts) if decision_parts else response_text
                span.set_attribute("message.response_content", decision_text)
            else:
                # Real text response — accumulate for merged agent_response at end_turn
                # Don't set role="assistant" so exporter skips conversation event
                span.set_attribute("message.role", "llm_response")
                if response_text:
                    span.set_attribute("message.content", response_text)
                    span.set_attribute("message.response_content", response_text)
                    self._turn_response_text = response_text  # Last response wins

            # Accumulate model name
            current_model = span.attributes.get("llm.model_name", "") if hasattr(span, 'attributes') else ""
            if not current_model:
                current_model = model
            if current_model and current_model not in self._turn_models:
                self._turn_models.append(current_model)

        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 = "AGENT_REQUEST" if is_subagent else "TOOL_REQUEST"

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

    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))
        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:
        # Don't emit events for chain spans — they're LangGraph internal routing.
        # BUT track the run_id → parent_run_id so child LLM/tool spans inside
        # sub-agents can traverse up to find their real parent span.
        key = str(run_id)
        self._chain_parents[key] = str(parent_run_id) if parent_run_id else None

    def on_chain_end(self, outputs: Any, *, run_id: Any, **kwargs: Any) -> None:
        self._chain_parents.pop(str(run_id), None)

    def on_chain_error(self, error: BaseException, *, run_id: Any, **kwargs: Any) -> None:
        self._chain_parents.pop(str(run_id), None)

    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
