"""
LangChain BaseCallbackHandler that emits plyra-trace spans.

Covers: chains, agents, tools, LLMs, retrievers.
Works with LangGraph via the same callback system.

Usage:
    from plyra_instrument_langchain import PlyraCallbackHandler

    handler = PlyraCallbackHandler()

    # LangChain
    chain = LLMChain(llm=llm, prompt=prompt, callbacks=[handler])

    # LangGraph
    app = workflow.compile()
    result = app.invoke(input, config={"callbacks": [handler]})
"""

from __future__ import annotations

import json
import logging
from typing import Any
from uuid import UUID

from langchain_core.callbacks.base import BaseCallbackHandler
from opentelemetry import context as otel_ctx
from opentelemetry import trace as trace_api

from plyra_trace.semconv import SpanAttributes, SpanKind

logger = logging.getLogger("plyra_instrument_langchain")


def _json_safe(v: Any) -> str:
    """Serialize value to JSON string, with fallback to str()."""
    try:
        return json.dumps(v, default=str)[:8000]
    except Exception:
        return str(v)[:8000]


class PlyraCallbackHandler(BaseCallbackHandler):
    """
    LangChain callback handler that emits plyra-trace spans.

    Uses run_id + parent_run_id from LangChain callbacks to reconstruct
    the parent-child tree correctly, even across LangGraph async boundaries.
    This replaces the old stack-based approach which lost context between nodes.
    """

    def __init__(self, tracer_provider: Any = None, session_id: str | None = None) -> None:
        if tracer_provider is not None:
            self._tracer = tracer_provider.get_tracer("plyra-instrument-langchain")
        else:
            self._tracer = trace_api.get_tracer("plyra-instrument-langchain")
        self.session_id = session_id
        # Map run_id (str) → Span
        self._run_spans: dict[str, Any] = {}
        # Map run_id (str) → OTel context token (for cleanup)
        self._run_tokens: dict[str, Any] = {}

    def _set_session(self, span: Any) -> None:
        """Attach session.id to span if set."""
        if self.session_id:
            span.set_attribute(SpanAttributes.SESSION_ID, self.session_id)

    def _start_span(
        self,
        name: str,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **attrs: Any,
    ) -> Any:
        """
        Start a span using LangChain's run_id / parent_run_id to reconstruct
        the correct parent-child relationship, even across async LangGraph nodes.

        If parent_run_id is known and its span still exists, the new span
        is created as a child. Otherwise, a root span starts a fresh trace.
        """
        run_key = str(run_id)
        parent_key = str(parent_run_id) if parent_run_id else None

        # Look up parent span from parent_run_id
        parent_span = self._run_spans.get(parent_key) if parent_key else None

        if parent_span is not None:
            ctx = trace_api.set_span_in_context(parent_span)
        else:
            # No parent known — start from the current ambient context
            ctx = otel_ctx.get_current()

        span = self._tracer.start_span(name, context=ctx)
        for k, v in attrs.items():
            if v is not None:
                span.set_attribute(k, v)

        self._run_spans[run_key] = span
        # Attach context so child spans in the same task can inherit it
        token = otel_ctx.attach(trace_api.set_span_in_context(span))
        self._run_tokens[run_key] = token
        return span

    def _end_span(self, run_id: UUID, output_attrs: dict[str, Any] | None = None) -> None:
        """End the span for this run_id and detach its context."""
        run_key = str(run_id)
        span = self._run_spans.pop(run_key, None)
        token = self._run_tokens.pop(run_key, None)
        if span:
            if output_attrs:
                for k, v in output_attrs.items():
                    if v is not None:
                        span.set_attribute(k, v)
            span.end()
        if token:
            otel_ctx.detach(token)

    # ── Chains ────────────────────────────────────────────────────────────────

    def on_chain_start(
        self,
        serialized: dict[str, Any] | None,
        inputs: dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        serialized = serialized or {}
        name = serialized.get("id", ["chain"])[-1] if serialized.get("id") else "chain"
        span = self._start_span(
            name,
            run_id,
            parent_run_id,
            **{
                SpanAttributes.SPAN_KIND: SpanKind.CHAIN,
                SpanAttributes.OPENINFERENCE_SPAN_KIND: "CHAIN",
                "langchain.span_type": "chain",
                SpanAttributes.INPUT_VALUE: _json_safe(inputs),
                SpanAttributes.INPUT_MIME_TYPE: "application/json",
            },
        )
        if metadata:
            span.set_attribute(SpanAttributes.METADATA, _json_safe(metadata))
        self._set_session(span)

    def on_chain_end(
        self,
        outputs: dict[str, Any],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        self._end_span(
            run_id,
            {
                SpanAttributes.OUTPUT_VALUE: _json_safe(outputs),
                SpanAttributes.OUTPUT_MIME_TYPE: "application/json",
            },
        )

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run_key = str(run_id)
        span = self._run_spans.get(run_key)
        if span:
            span.set_status(trace_api.StatusCode.ERROR, str(error))
        self._end_span(run_id)

    # ── LLMs ──────────────────────────────────────────────────────────────────

    def on_llm_start(
        self,
        serialized: dict[str, Any] | None,
        prompts: list[str],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        serialized = serialized or {}
        invoc = kwargs.get("invocation_params", {}) or {}
        model = invoc.get("model_name") or invoc.get("model") or serialized.get("kwargs", {}).get("model_name") or "llm"
        name = f"llm.{model}" if model and model != "llm" else "llm"
        span = self._start_span(
            name,
            run_id,
            parent_run_id,
            **{
                SpanAttributes.SPAN_KIND: SpanKind.LLM,
                SpanAttributes.OPENINFERENCE_SPAN_KIND: "LLM",
                "langchain.span_type": "llm",
                SpanAttributes.LLM_MODEL_NAME: model,
                SpanAttributes.AGENT_FRAMEWORK: "langchain",
            },
        )
        if prompts:
            span.set_attribute(SpanAttributes.INPUT_VALUE, prompts[0][:4000])
        if invoc:
            params = {k: v for k, v in invoc.items() if k in ("temperature", "max_tokens", "top_p") and v is not None}
            if params:
                span.set_attribute(SpanAttributes.LLM_INVOCATION_PARAMETERS, json.dumps(params))
        self._set_session(span)

    def on_llm_end(
        self,
        response: Any,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        attrs: dict[str, Any] = {}
        try:
            if hasattr(response, "generations") and response.generations:
                text = response.generations[0][0].text if response.generations[0] else ""
                attrs[SpanAttributes.OUTPUT_VALUE] = text[:4000]
            if hasattr(response, "llm_output") and response.llm_output:
                usage = response.llm_output.get("token_usage") or {}
                prompt_toks = usage.get("prompt_tokens", 0) or 0
                completion_toks = usage.get("completion_tokens", 0) or 0
                if prompt_toks or completion_toks:
                    attrs[SpanAttributes.LLM_TOKEN_COUNT_PROMPT] = prompt_toks
                    attrs[SpanAttributes.LLM_TOKEN_COUNT_COMPLETION] = completion_toks
                    attrs[SpanAttributes.LLM_TOKEN_COUNT_TOTAL] = prompt_toks + completion_toks
        except Exception:  # noqa: BLE001
            pass
        self._end_span(run_id, attrs)

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run_key = str(run_id)
        span = self._run_spans.get(run_key)
        if span:
            span.set_status(trace_api.StatusCode.ERROR, str(error))
        self._end_span(run_id)

    # ── Tools ─────────────────────────────────────────────────────────────────

    def on_tool_start(
        self,
        serialized: dict[str, Any] | None,
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        serialized = serialized or {}
        name = serialized.get("name") or "tool"
        span = self._start_span(
            name,
            run_id,
            parent_run_id,
            **{
                SpanAttributes.SPAN_KIND: SpanKind.TOOL,
                SpanAttributes.OPENINFERENCE_SPAN_KIND: "TOOL",
                "langchain.span_type": "tool",
                SpanAttributes.TOOL_NAME: name,
                SpanAttributes.INPUT_VALUE: str(input_str)[:4000],
            },
        )
        if serialized.get("description"):
            span.set_attribute(SpanAttributes.TOOL_DESCRIPTION, serialized["description"])
        self._set_session(span)

    def on_tool_end(
        self,
        output: str,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        self._end_span(run_id, {SpanAttributes.OUTPUT_VALUE: str(output)[:4000]})

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run_key = str(run_id)
        span = self._run_spans.get(run_key)
        if span:
            span.set_status(trace_api.StatusCode.ERROR, str(error))
        self._end_span(run_id)

    # ── Agents ────────────────────────────────────────────────────────────────

    def on_agent_action(
        self,
        action: Any,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        tool_name = getattr(action, "tool", "agent_action")
        span = self._start_span(
            f"agent.{tool_name}",
            run_id,
            parent_run_id,
            **{
                SpanAttributes.SPAN_KIND: SpanKind.AGENT,
                SpanAttributes.OPENINFERENCE_SPAN_KIND: "AGENT",
                "langchain.span_type": "agent",
                SpanAttributes.AGENT_FRAMEWORK: "langchain",
            },
        )
        log = getattr(action, "log", "")
        if log:
            span.set_attribute(SpanAttributes.INPUT_VALUE, str(log)[:4000])
        self._set_session(span)

    def on_agent_finish(
        self,
        finish: Any,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        output = getattr(finish, "return_values", {})
        self._end_span(run_id, {SpanAttributes.OUTPUT_VALUE: _json_safe(output)})

    # ── Retrievers ────────────────────────────────────────────────────────────

    def on_retriever_start(
        self,
        serialized: dict[str, Any] | None,
        query: str,
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        **kwargs: Any,
    ) -> None:
        span = self._start_span(
            "retriever",
            run_id,
            parent_run_id,
            **{
                SpanAttributes.SPAN_KIND: SpanKind.RETRIEVER,
                SpanAttributes.OPENINFERENCE_SPAN_KIND: "RETRIEVER",
                "langchain.span_type": "retriever",
                SpanAttributes.RETRIEVAL_QUERY: query[:2000],
            },
        )
        self._set_session(span)

    def on_retriever_end(
        self,
        documents: list[Any],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        attrs: dict[str, Any] = {"retrieval.documents.count": len(documents)}
        if documents:
            docs_preview = [{"content": str(getattr(d, "page_content", d))[:500]} for d in documents[:5]]
            attrs[SpanAttributes.OUTPUT_VALUE] = _json_safe(docs_preview)
        self._end_span(run_id, attrs)

    def on_retriever_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        run_key = str(run_id)
        span = self._run_spans.get(run_key)
        if span:
            span.set_status(trace_api.StatusCode.ERROR, str(error))
        self._end_span(run_id)
