"""
Cascade SDK - LangGraph / LangChain Auto-Instrumentation

Automatically traces all LangGraph agent executions, LLM calls,
tool calls, and chain runs with zero code changes.

Two APIs are provided:

1. **setup_langgraph()** / **instrument_langgraph()** — instruments all
   LangGraph/LangChain runs globally (one call).
2. **CascadeLangChainHandler** — a callback handler you can pass
   manually via ``config={"callbacks": [CascadeLangChainHandler()]}``.

Usage::

    from cascade import init_tracing
    from cascade.integrations import setup_langgraph

    init_tracing(project="my_langgraph_app")
    setup_langgraph()

    # Existing LangGraph code -- no changes needed
    app = graph.compile()
    result = app.invoke({"messages": [...]})
"""

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

from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode, SpanKind

from cascade.tracing import get_tracer, get_current_agent

logger = logging.getLogger(__name__)

_instrumented = False


def instrument_langgraph() -> None:
    """
    Auto-instrument LangGraph/LangChain with Cascade tracing.

    Call this once after ``init_tracing()`` and before running any
    LangGraph code. All subsequent agent invocations will be traced
    automatically.

    Requires: ``pip install langgraph langchain-core``

    Raises:
        ImportError: If langchain-core is not installed.
    """
    global _instrumented
    if _instrumented:
        logger.debug("LangGraph already instrumented, skipping")
        return

    try:
        from langchain_core.callbacks import BaseCallbackHandler
        from langchain_core.callbacks.manager import CallbackManager
    except ImportError:
        raise ImportError(
            "LangGraph/LangChain not found. Install with: "
            "pip install langgraph langchain-core"
        )

    handler = CascadeLangChainHandler()

    # Patch the default callback manager to include our handler
    # so every LangChain/LangGraph invocation is traced automatically.
    try:
        from langchain_core.callbacks.manager import (
            CallbackManager,
            AsyncCallbackManager,
        )

        _original_configure = CallbackManager.configure
        _original_async_configure = AsyncCallbackManager.configure

        @classmethod  # type: ignore[misc]
        def patched_configure(cls, *args, **kwargs):
            manager = _original_configure.__func__(cls, *args, **kwargs)
            # Add our handler if not already present
            if not any(isinstance(h, CascadeLangChainHandler) for h in manager.handlers):
                manager.add_handler(handler)
            return manager

        @classmethod  # type: ignore[misc]
        def patched_async_configure(cls, *args, **kwargs):
            manager = _original_async_configure.__func__(cls, *args, **kwargs)
            if not any(isinstance(h, CascadeLangChainHandler) for h in manager.handlers):
                manager.add_handler(handler)
            return manager

        CallbackManager.configure = patched_configure
        AsyncCallbackManager.configure = patched_async_configure

    except Exception as e:
        logger.warning(f"Could not auto-patch CallbackManager: {e}")
        logger.info(
            "Pass the handler manually: "
            "app.invoke(..., config={'callbacks': [CascadeLangChainHandler()]})"
        )

    _instrumented = True
    logger.info("Cascade: LangGraph/LangChain auto-instrumentation enabled")


# Convenience alias matching the setup_pydantic_ai() naming convention.
setup_langgraph = instrument_langgraph


def _safe_str(value: Any, max_len: int = 4000) -> str:
    """Safely convert a value to a truncated string."""
    try:
        if isinstance(value, str):
            s = value
        elif isinstance(value, dict) or isinstance(value, list):
            s = json.dumps(value, default=str, ensure_ascii=False)
        else:
            s = str(value)
        return s[:max_len] if len(s) > max_len else s
    except Exception:
        return "<unserializable>"


class CascadeLangChainHandler:
    """
    LangChain BaseCallbackHandler that translates events into Cascade spans.

    Uses LangChain's ``parent_run_id`` to build the span hierarchy instead
    of ``context_api.attach/detach``, which breaks across async boundaries
    (``ainvoke``).
    """

    def __init__(self):
        # run_id → span (no context tokens needed)
        self._spans: Dict[str, Any] = {}

        from langchain_core.callbacks import BaseCallbackHandler

        self.__class__ = type(
            "CascadeLangChainHandler",
            (BaseCallbackHandler,),
            {k: v for k, v in CascadeLangChainHandler.__dict__.items() if not k.startswith("__")},
        )

    @property
    def always_verbose(self) -> bool:
        return True

    @property
    def raise_error(self) -> bool:
        return False

    # Names that are graph routing/plumbing — skip creating spans for these.
    _SKIP_NAMES = frozenset({
        "should_continue", "route", "router", "routing",
        "decide", "decide_next", "check_condition",
        "__start__", "__end__",
        "ChannelWrite", "ChannelRead",
        "RunnableSequence", "RunnableLambda",
        "RunnableParallel", "RunnablePassthrough",
    })

    def _parent_context(self, parent_run_id: Optional[UUID]):
        """Return an OTel context with the parent span set, or None."""
        if parent_run_id:
            parent = self._spans.get(str(parent_run_id))
            if parent:
                return trace.set_span_in_context(parent["span"])
        return None

    def _should_skip(self, name: str) -> bool:
        """Return True if this chain name is routing plumbing we should hide."""
        if name in self._SKIP_NAMES:
            return True
        lower = name.lower()
        if lower.startswith("runnable") or lower.startswith("channel"):
            return True
        return False

    # ------------------------------------------------------------------
    # Chain (Agent / Graph node / Runnable)
    # ------------------------------------------------------------------

    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:
        try:
            tracer = get_tracer()
            if not tracer:
                return

            serialized = serialized or {}
            name = serialized.get("name") or (serialized.get("id", ["Unknown"]) or ["Unknown"])[-1]
            if name == "Unknown":
                name = kwargs.get("name", name)

            if self._should_skip(name):
                return

            name_lower = name.lower()
            is_agent = "agent" in name_lower or (tags and "agent" in tags)
            span_type = "agent" if is_agent else "function"

            ctx = self._parent_context(parent_run_id)
            span = tracer.start_span(name, kind=SpanKind.INTERNAL, context=ctx)
            span.set_attribute("cascade.span_type", span_type)

            if is_agent:
                span.set_attribute("cascade.agent_name", name)
            if span_type == "function":
                span.set_attribute("function.name", name)
                span.set_attribute("function.input", _safe_str(inputs))

            self._spans[str(run_id)] = {"span": span, "name": name, "type": span_type}
        except Exception as e:
            logger.debug(f"Cascade: on_chain_start tracing failed: {e}")

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)
        # If this run_id was skipped in on_chain_start, ignore the end too.
        if rid not in self._spans:
            return
        try:
            entry = self._spans.pop(rid)
            span = entry["span"]
            if entry["type"] == "function":
                span.set_attribute("function.output", _safe_str(outputs))
            span.set_status(Status(StatusCode.OK))
            span.end()
        except Exception as e:
            logger.debug(f"Cascade: on_chain_end tracing failed: {e}")

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)
        if rid not in self._spans:
            return
        try:
            entry = self._spans.pop(rid)
            span = entry["span"]
            span.set_status(Status(StatusCode.ERROR, str(error)))
            span.record_exception(error)
            span.end()
        except Exception as e:
            logger.debug(f"Cascade: on_chain_error tracing failed: {e}")

    # ------------------------------------------------------------------
    # LLM
    # ------------------------------------------------------------------

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            tracer = get_tracer()
            if not tracer:
                return

            serialized = serialized or {}
            model = serialized.get("kwargs", {}).get("model_name") or (serialized.get("id", ["unknown"]) or ["unknown"])[-1]
            name = f"llm.{model}"

            ctx = self._parent_context(parent_run_id)
            span = tracer.start_span(name, kind=SpanKind.INTERNAL, context=ctx)
            span.set_attribute("cascade.span_type", "llm")
            span.set_attribute("llm.model", model)
            span.set_attribute("llm.provider", (serialized.get("id", ["unknown"]) or ["unknown"])[0])

            prompt_text = "\n".join(prompts) if prompts else ""
            span.set_attribute("llm.prompt", _safe_str(prompt_text))

            current_agent = get_current_agent()
            if current_agent:
                span.set_attribute("cascade.agent_name", current_agent)

            self._spans[str(run_id)] = {"span": span, "name": name, "type": "llm"}
        except Exception as e:
            logger.debug(f"Cascade: on_llm_start tracing failed: {e}")

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[Any]],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            tracer = get_tracer()
            if not tracer:
                return

            serialized = serialized or {}
            model = (
                serialized.get("kwargs", {}).get("model_name")
                or serialized.get("kwargs", {}).get("model")
                or (serialized.get("id", ["unknown"]) or ["unknown"])[-1]
            )
            name = f"llm.{model}"

            ctx = self._parent_context(parent_run_id)
            span = tracer.start_span(name, kind=SpanKind.INTERNAL, context=ctx)
            span.set_attribute("cascade.span_type", "llm")
            span.set_attribute("llm.model", model)

            id_parts = serialized.get("id", [])
            provider = "unknown"
            for part in id_parts:
                if "openai" in str(part).lower():
                    provider = "openai"
                    break
                elif "anthropic" in str(part).lower():
                    provider = "anthropic"
                    break
            span.set_attribute("llm.provider", provider)

            prompt_parts = []
            for msg_list in messages:
                for msg in msg_list:
                    role = getattr(msg, "type", "unknown")
                    content = getattr(msg, "content", str(msg))
                    prompt_parts.append(f"[{role}] {content}")
            span.set_attribute("llm.prompt", _safe_str("\n".join(prompt_parts)))

            current_agent = get_current_agent()
            if current_agent:
                span.set_attribute("cascade.agent_name", current_agent)

            self._spans[str(run_id)] = {"span": span, "name": name, "type": "llm"}
        except Exception as e:
            logger.debug(f"Cascade: on_chat_model_start tracing failed: {e}")

    def on_llm_end(
        self,
        response: Any,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        try:
            entry = self._spans.pop(str(run_id), None)
            if not entry:
                return
            span = entry["span"]

            try:
                generations = response.generations
                if generations and generations[0]:
                    gen = generations[0][0]
                    msg = getattr(gen, "message", None)
                    completion = getattr(gen, "text", "") or (getattr(msg, "content", "") if msg else "")
                    span.set_attribute("llm.completion", _safe_str(completion))

                    token_usage = None
                    if msg:
                        umeta = getattr(msg, "usage_metadata", None)
                        if umeta and isinstance(umeta, dict):
                            token_usage = umeta
                        elif hasattr(msg, "response_metadata"):
                            rmeta = getattr(msg, "response_metadata", {}) or {}
                            token_usage = rmeta.get("token_usage") or rmeta.get("usage")

                    if not token_usage:
                        llm_output = getattr(response, "llm_output", {}) or {}
                        token_usage = llm_output.get("token_usage", {})

                    if token_usage:
                        in_tok = token_usage.get("input_tokens") or token_usage.get("prompt_tokens") or 0
                        out_tok = token_usage.get("output_tokens") or token_usage.get("completion_tokens") or 0
                        total = token_usage.get("total_tokens") or (in_tok + out_tok)
                        span.set_attribute("llm.input_tokens", in_tok)
                        span.set_attribute("llm.output_tokens", out_tok)
                        span.set_attribute("llm.total_tokens", total)
            except Exception as e:
                logger.debug(f"Error extracting LLM response: {e}")

            span.set_status(Status(StatusCode.OK))
            span.end()
        except Exception as e:
            logger.debug(f"Cascade: on_llm_end tracing failed: {e}")

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        try:
            entry = self._spans.pop(str(run_id), None)
            if not entry:
                return
            span = entry["span"]
            span.set_status(Status(StatusCode.ERROR, str(error)))
            span.record_exception(error)
            span.end()
        except Exception as e:
            logger.debug(f"Cascade: on_llm_error tracing failed: {e}")

    # ------------------------------------------------------------------
    # Tool
    # ------------------------------------------------------------------

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            tracer = get_tracer()
            if not tracer:
                return

            serialized = serialized or {}
            tool_name = serialized.get("name") or (serialized.get("id", ["unknown"]) or ["unknown"])[-1]

            ctx = self._parent_context(parent_run_id)
            span = tracer.start_span(tool_name, kind=SpanKind.INTERNAL, context=ctx)
            span.set_attribute("cascade.span_type", "tool")
            span.set_attribute("tool.name", tool_name)
            span.set_attribute("tool.input", _safe_str(input_str))

            if serialized.get("description"):
                span.set_attribute("tool.description", serialized["description"][:200])

            current_agent = get_current_agent()
            if current_agent:
                span.set_attribute("cascade.agent_name", current_agent)

            self._spans[str(run_id)] = {"span": span, "name": tool_name, "type": "tool"}
        except Exception as e:
            logger.debug(f"Cascade: on_tool_start tracing failed: {e}")

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        try:
            entry = self._spans.pop(str(run_id), None)
            if not entry:
                return
            span = entry["span"]
            span.set_attribute("tool.output", _safe_str(output))
            span.set_status(Status(StatusCode.OK))
            span.end()
        except Exception as e:
            logger.debug(f"Cascade: on_tool_end tracing failed: {e}")

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        try:
            entry = self._spans.pop(str(run_id), None)
            if not entry:
                return
            span = entry["span"]
            span.set_attribute("tool.error", str(error))
            span.set_status(Status(StatusCode.ERROR, str(error)))
            span.record_exception(error)
            span.end()
        except Exception as e:
            logger.debug(f"Cascade: on_tool_error tracing failed: {e}")

    # ------------------------------------------------------------------
    # Retriever (maps to tool spans)
    # ------------------------------------------------------------------

    def on_retriever_start(
        self,
        serialized: Dict[str, Any],
        query: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        try:
            tracer = get_tracer()
            if not tracer:
                return

            serialized = serialized or {}
            name = serialized.get("name", "retriever")

            ctx = self._parent_context(parent_run_id)
            span = tracer.start_span(name, kind=SpanKind.INTERNAL, context=ctx)
            span.set_attribute("cascade.span_type", "tool")
            span.set_attribute("tool.name", name)
            span.set_attribute("tool.input", _safe_str(query))

            self._spans[str(run_id)] = {"span": span, "name": name, "type": "tool"}
        except Exception as e:
            logger.debug(f"Cascade: on_retriever_start tracing failed: {e}")

    def on_retriever_end(
        self,
        documents: List[Any],
        *,
        run_id: UUID,
        **kwargs: Any,
    ) -> None:
        try:
            entry = self._spans.pop(str(run_id), None)
            if not entry:
                return
            span = entry["span"]
            doc_texts = []
            for doc in documents[:10]:
                content = getattr(doc, "page_content", str(doc))
                doc_texts.append(content[:500])
            span.set_attribute("tool.output", _safe_str("\n---\n".join(doc_texts)))
            span.set_status(Status(StatusCode.OK))
            span.end()
        except Exception as e:
            logger.debug(f"Cascade: on_retriever_end tracing failed: {e}")
