"""
LangChain callback handler for Risicare tracing.

Implements langchain_core.callbacks.BaseCallbackHandler to capture
LangChain execution events as Risicare spans with proper parent-child
linking, content capture, and error recording.

Span Mapping:
    on_chain_start/end       → INTERNAL span (langchain.chain/{class})
    on_llm_start/end         → LLM_CALL span (langchain.llm/{model}), phase=THINK
    on_chat_model_start/end  → LLM_CALL span (langchain.chat/{model}), phase=THINK
    on_tool_start/end        → TOOL_CALL span (langchain.tool/{name}), phase=ACT
    on_retriever_start/end   → RETRIEVAL span (langchain.retriever/{name}), phase=OBSERVE
    on_agent_action          → DECISION span (langchain.agent_action/{tool}), phase=DECIDE
    on_agent_finish          → sets result on parent span
"""

from __future__ import annotations

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

from risicare.integrations._base import (
    get_tracer,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
    create_framework_attributes,
    get_framework_version,
)

logger = logging.getLogger(__name__)


class _SpanEntry:
    """Tracks a span and its metadata during callback lifecycle."""

    __slots__ = ("span", "token", "start_time")

    def __init__(self, span: Any, token: Any, start_time: float) -> None:
        self.span = span
        self.token = token
        self.start_time = start_time


class RisicareCallbackHandler:
    """
    LangChain callback handler that creates Risicare spans.

    Maps LangChain execution events (chain, llm, tool, retriever, agent)
    to Risicare spans with proper parent-child linking via run_id/parent_run_id.

    Can be used in two ways:
    1. Automatically injected by the LangChain patches (zero config)
    2. Manually passed to chain.invoke() config

    Example:
        from risicare.integrations.langchain import RisicareCallbackHandler
        chain.invoke("hello", config={"callbacks": [RisicareCallbackHandler()]})
    """

    # Maximum number of tracked in-flight spans. If exceeded, oldest entries
    # are evicted to prevent unbounded memory growth when on_chain_end never
    # fires (e.g. LangChain internal crash, unhandled exception).
    _MAX_RUN_SPANS = 10_000

    # Maximum age (seconds) before a span entry is considered stale and
    # eligible for time-based eviction.  30 minutes is generous — even the
    # longest LangChain chains complete in under 10 minutes.  (P2-9)
    _SPAN_TTL_SECONDS = 1800

    def __init__(self) -> None:
        self._run_spans: Dict[UUID, _SpanEntry] = {}
        self._trace_content = should_trace_content()
        self._framework_version = get_framework_version("langchain-core")

    # -------------------------------------------------------------------------
    # Internal helpers
    # -------------------------------------------------------------------------

    def _start_span(
        self,
        run_id: UUID,
        parent_run_id: Optional[UUID],
        name: str,
        kind_str: str,
        phase_str: Optional[str] = None,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> Optional[Any]:
        """Start a span and register it by run_id."""
        try:
            tracer = get_tracer()
            if tracer is None or not tracer.is_enabled:
                return None

            from risicare_core import SpanKind, SemanticPhase

            kind_map = {
                "INTERNAL": SpanKind.INTERNAL,
                "LLM_CALL": SpanKind.LLM_CALL,
                "TOOL_CALL": SpanKind.TOOL_CALL,
                "RETRIEVAL": SpanKind.RETRIEVAL,
                "DECISION": SpanKind.DECISION,
            }
            kind = kind_map.get(kind_str, SpanKind.INTERNAL)

            # Build attributes
            span_attrs = create_framework_attributes(
                "langchain", self._framework_version
            )
            span_attrs["framework.langchain.run_id"] = str(run_id)
            if parent_run_id:
                span_attrs["framework.langchain.parent_run_id"] = str(parent_run_id)
            if attributes:
                span_attrs.update(attributes)

            # Resolve parent span from run_id
            parent_entry = self._run_spans.get(parent_run_id) if parent_run_id else None

            # Create span with manual lifecycle (no context manager)
            # We need to use start_span_no_context because callback events
            # arrive asynchronously — we can't use a with statement
            parent_span_id = parent_entry.span.span_id if parent_entry else None
            trace_id = parent_entry.span.trace_id if parent_entry else None

            span = tracer.start_span_no_context(
                name=name,
                kind=kind,
                attributes=span_attrs,
                parent_span_id=parent_span_id,
                trace_id=trace_id,
            )

            # Set semantic phase
            if phase_str:
                phase_map = {
                    "THINK": SemanticPhase.THINK,
                    "ACT": SemanticPhase.ACT,
                    "OBSERVE": SemanticPhase.OBSERVE,
                    "DECIDE": SemanticPhase.DECIDE,
                }
                phase = phase_map.get(phase_str)
                if phase:
                    span.semantic_phase = phase

            # Guard against unbounded growth: if _run_spans exceeds the
            # limit, first try time-based eviction (P2-9), then fall back
            # to FIFO eviction of the oldest 10%.
            if len(self._run_spans) >= self._MAX_RUN_SPANS:
                now_mono = time.perf_counter()
                stale_ids = [
                    rid
                    for rid, entry in self._run_spans.items()
                    if (now_mono - entry.start_time) > self._SPAN_TTL_SECONDS
                ]
                if stale_ids:
                    for rid in stale_ids:
                        del self._run_spans[rid]
                    logger.warning(
                        "Evicted %d stale span entries (TTL=%ds) from "
                        "callback handler. This may indicate LangChain "
                        "callbacks are not firing on_chain_end.",
                        len(stale_ids),
                        self._SPAN_TTL_SECONDS,
                    )
                else:
                    # No stale entries — fall back to FIFO eviction
                    evict_count = max(1, self._MAX_RUN_SPANS // 10)
                    evict_ids = list(self._run_spans.keys())[:evict_count]
                    for evict_id in evict_ids:
                        del self._run_spans[evict_id]
                    logger.warning(
                        "Evicted %d orphaned span entries from callback "
                        "handler (limit: %d, no stale entries found). "
                        "This may indicate a burst of concurrent spans.",
                        evict_count,
                        self._MAX_RUN_SPANS,
                    )

            entry = _SpanEntry(
                span=span,
                token=None,  # No context token — manual lifecycle
                start_time=time.perf_counter(),
            )
            self._run_spans[run_id] = entry
            return span

        except Exception:
            logger.debug("Failed to start LangChain span", exc_info=True)
            return None

    def _end_span(
        self,
        run_id: UUID,
        error: Optional[BaseException] = None,
        attributes: Optional[Dict[str, Any]] = None,
    ) -> None:
        """End a span by run_id."""
        entry = self._run_spans.pop(run_id, None)
        if entry is None:
            return

        try:
            span = entry.span
            latency_ms = (time.perf_counter() - entry.start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if attributes:
                for k, v in attributes.items():
                    safe_set_attribute(span, k, v)

            if error:
                record_error(span, error)
            else:
                from risicare_core import SpanStatus

                if span.status == SpanStatus.UNSET:
                    span.mark_ok()

            span.end()

            # Export
            tracer = get_tracer()
            if tracer:
                tracer._export_span(span)

        except Exception:
            logger.debug("Failed to end LangChain span", exc_info=True)

    # -------------------------------------------------------------------------
    # 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."""
        class_name = serialized.get("name") or serialized.get("id", ["unknown"])[-1]
        attrs: Dict[str, Any] = {
            "framework.langchain.chain_type": class_name,
        }
        if tags:
            attrs["framework.langchain.tags"] = tags
        if self._trace_content and inputs:
            input_str = str(inputs)
            attrs["gen_ai.prompt.content"] = truncate_content(input_str)

        self._start_span(
            run_id=run_id,
            parent_run_id=parent_run_id,
            name=f"langchain.chain/{class_name}",
            kind_str="INTERNAL",
            attributes=attrs,
        )

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a chain finishes."""
        attrs: Dict[str, Any] = {}
        if self._trace_content and outputs:
            output_str = str(outputs)
            attrs["gen_ai.completion.content"] = truncate_content(output_str)
        self._end_span(run_id, attributes=attrs)

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a chain errors."""
        self._end_span(run_id, 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,
        invocation_params: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an LLM starts (text completion models)."""
        model = "unknown"
        if invocation_params:
            model = invocation_params.get("model_name") or invocation_params.get(
                "model", "unknown"
            )

        attrs: Dict[str, Any] = {
            "gen_ai.system": "langchain",
            "gen_ai.request.model": model,
        }
        if self._trace_content:
            attrs["gen_ai.prompt.count"] = len(prompts)
            if prompts:
                attrs["gen_ai.prompt.content"] = truncate_content(prompts[0])

        self._start_span(
            run_id=run_id,
            parent_run_id=parent_run_id,
            name=f"langchain.llm/{model}",
            kind_str="LLM_CALL",
            phase_str="THINK",
            attributes=attrs,
        )

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[Any]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        invocation_params: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a chat model starts (ChatOpenAI, ChatAnthropic, etc.)."""
        model = "unknown"
        if invocation_params:
            model = invocation_params.get("model_name") or invocation_params.get(
                "model", "unknown"
            )

        attrs: Dict[str, Any] = {
            "gen_ai.system": "langchain",
            "gen_ai.request.model": model,
        }

        if messages:
            # messages is List[List[BaseMessage]] — outer list is per-invocation,
            # inner list is the message sequence. Guard against empty or
            # non-list first element (varies by LangChain version).
            first_batch = messages[0] if messages else []
            if not isinstance(first_batch, (list, tuple)):
                first_batch = [first_batch] if first_batch is not None else []
            flat_messages = first_batch
            attrs["gen_ai.prompt.count"] = len(flat_messages)

            if self._trace_content:
                for i, msg in enumerate(flat_messages[:10]):
                    role = getattr(msg, "type", "unknown")
                    content = getattr(msg, "content", "")
                    attrs[f"gen_ai.prompt.{i}.role"] = role
                    if isinstance(content, str):
                        attrs[f"gen_ai.prompt.{i}.content"] = truncate_content(content)

        self._start_span(
            run_id=run_id,
            parent_run_id=parent_run_id,
            name=f"langchain.chat/{model}",
            kind_str="LLM_CALL",
            phase_str="THINK",
            attributes=attrs,
        )

    def on_llm_end(
        self,
        response: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an LLM finishes."""
        attrs: Dict[str, Any] = {}

        # Extract token usage from LLMResult
        llm_output = getattr(response, "llm_output", None) or {}
        token_usage = llm_output.get("token_usage") or {}
        if token_usage:
            attrs["gen_ai.usage.prompt_tokens"] = token_usage.get("prompt_tokens", 0)
            attrs["gen_ai.usage.completion_tokens"] = token_usage.get(
                "completion_tokens", 0
            )
            attrs["gen_ai.usage.total_tokens"] = token_usage.get("total_tokens", 0)

        # Extract model from LLMResult
        model = llm_output.get("model_name")
        if model:
            attrs["gen_ai.response.model"] = model

        # Extract completion content
        generations = getattr(response, "generations", [])
        if generations and self._trace_content:
            first_gen = generations[0]
            if first_gen:
                gen = first_gen[0] if isinstance(first_gen, list) else first_gen
                text = getattr(gen, "text", "")
                if text:
                    attrs["gen_ai.completion.content"] = truncate_content(text)

                # Check for message-based generations
                message = getattr(gen, "message", None)
                if message:
                    content = getattr(message, "content", "")
                    if content and isinstance(content, str):
                        attrs["gen_ai.completion.content"] = truncate_content(content)

                    # Tool calls
                    tool_calls = getattr(message, "tool_calls", None)
                    if tool_calls:
                        attrs["gen_ai.completion.tool_calls"] = len(tool_calls)
                        for j, tc in enumerate(tool_calls[:5]):
                            tc_name = (
                                tc.get("name", "")
                                if isinstance(tc, dict)
                                else getattr(tc, "name", "")
                            )
                            attrs[f"gen_ai.completion.tool_call.{j}.name"] = tc_name

        self._end_span(run_id, attributes=attrs)

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an LLM errors."""
        self._end_span(run_id, 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,
        **kwargs: Any,
    ) -> None:
        """Called when a tool starts."""
        tool_name = serialized.get("name", "unknown")
        attrs: Dict[str, Any] = {
            "tool.name": tool_name,
        }
        if self._trace_content:
            attrs["tool.input"] = truncate_content(input_str)
        else:
            attrs["tool.input_length"] = len(input_str) if input_str else 0

        self._start_span(
            run_id=run_id,
            parent_run_id=parent_run_id,
            name=f"langchain.tool/{tool_name}",
            kind_str="TOOL_CALL",
            phase_str="ACT",
            attributes=attrs,
        )

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a tool finishes."""
        attrs: Dict[str, Any] = {}
        output_str = str(output) if output is not None else ""
        if self._trace_content:
            attrs["tool.output"] = truncate_content(output_str)
        else:
            attrs["tool.output_length"] = len(output_str)

        self._end_span(run_id, attributes=attrs)

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a tool errors."""
        self._end_span(run_id, error=error)

    # -------------------------------------------------------------------------
    # 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 a retriever starts."""
        retriever_name = serialized.get("name") or serialized.get(
            "id", ["unknown"]
        )[-1]
        attrs: Dict[str, Any] = {
            "retriever.name": retriever_name,
        }
        if self._trace_content:
            attrs["retriever.query"] = truncate_content(query)

        self._start_span(
            run_id=run_id,
            parent_run_id=parent_run_id,
            name=f"langchain.retriever/{retriever_name}",
            kind_str="RETRIEVAL",
            phase_str="OBSERVE",
            attributes=attrs,
        )

    def on_retriever_end(
        self,
        documents: Sequence[Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a retriever finishes."""
        attrs: Dict[str, Any] = {
            "retriever.document_count": len(documents),
        }
        if self._trace_content and documents:
            # Capture first doc content preview
            first_doc = documents[0]
            content = getattr(first_doc, "page_content", str(first_doc))
            if isinstance(content, str):
                attrs["retriever.first_doc"] = truncate_content(content, 2000)

        self._end_span(run_id, attributes=attrs)

    def on_retriever_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a retriever errors."""
        self._end_span(run_id, 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 an agent decides on an action."""
        tool = getattr(action, "tool", "unknown")
        tool_input = getattr(action, "tool_input", "")
        log = getattr(action, "log", "")

        attrs: Dict[str, Any] = {
            "agent.action.tool": tool,
        }
        if self._trace_content:
            if isinstance(tool_input, str):
                attrs["agent.action.input"] = truncate_content(tool_input)
            elif isinstance(tool_input, dict):
                attrs["agent.action.input"] = truncate_content(str(tool_input))
            if log:
                attrs["agent.action.log"] = truncate_content(log, 5000)

        self._start_span(
            run_id=run_id,
            parent_run_id=parent_run_id,
            name=f"langchain.agent_action/{tool}",
            kind_str="DECISION",
            phase_str="DECIDE",
            attributes=attrs,
        )

    def on_agent_finish(
        self,
        finish: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when an agent finishes."""
        # Agent finish doesn't always have a matching on_agent_action start,
        # so we set result on the parent span if it exists
        parent_entry = self._run_spans.get(parent_run_id) if parent_run_id else None
        if parent_entry and self._trace_content:
            return_values = getattr(finish, "return_values", {})
            output = return_values.get("output", "")
            if output:
                safe_set_attribute(
                    parent_entry.span,
                    "agent.finish.output",
                    truncate_content(str(output)),
                )

        # End the action span if one exists for this run_id
        entry = self._run_spans.get(run_id)
        if entry:
            self._end_span(run_id)

    # -------------------------------------------------------------------------
    # Text callbacks (no-ops — text generation is captured by on_llm_*)
    # -------------------------------------------------------------------------

    def on_text(self, text: str, **kwargs: Any) -> None:
        """Called with text output (no-op)."""
        pass

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Called on streaming tokens (no-op — we capture full response)."""
        pass
