"""
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

    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 ""

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

    def _start_span(self, run_id: str, name: str, span_type: str, attrs: Optional[Dict] = None):
        span = self._tracer.start_span(name)
        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 _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()
        context.detach(token)

    # ---- LLM callbacks ----

    def _format_messages_as_prompt(self, messages: List[Any]) -> str:
        """Format a list of LangChain messages into a readable prompt string.

        Captures system prompt + conversation context, not just the last message.
        This becomes agentTraceContent on the agent_trace event.
        """
        if not messages:
            return ""
        # messages[0] is the list of messages for the first batch
        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
            # Truncate individual messages but keep the structure
            truncated = content[:300] if len(content) > 300 else content
            parts.append(f"[{role}] {truncated}")
        # Join all messages, truncate total to 500 chars
        full = "\n".join(parts)
        return full[:500] if len(full) > 500 else full

    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:
        model = serialized.get("kwargs", {}).get("model_name", "") or serialized.get("id", [""])[-1]
        # Capture full prompt context, not just first message
        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",
        })

    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:
        model = kwargs.get("invocation_params", {}).get("model_name", "") or serialized.get("kwargs", {}).get("model_name", "") or serialized.get("id", [""])[-1]
        # Full prompt context for agentTraceContent (system + context + user)
        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
        self._start_span(run_id, f"llm.{model}", "GENERATION", {
            "llm.model_name": model,
            "message.content": user_content or full_prompt[:200],
            "message.trace_content": full_prompt,
            "message.role": "user",
        })

    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)
        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.get("name", "") or serialized.get("id", [""])[-1]
        self._start_span(run_id, f"tool.{tool_name}", "TOOL", {
            "tool.name": tool_name,
            "tool.input": str(input_str)[:500],
        })

    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 = serialized.get("name", "") or serialized.get("id", [""])[-1]
        self._start_span(run_id, f"chain.{chain_name}", "CHAIN")

    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
