"""LangGraph callback handler for Arzule observability.

Since LangGraph is built on LangChain and uses BaseCallbackHandler,
this handler extends LangChain's callback system to capture LangGraph-specific events.
"""

from __future__ import annotations

import sys
import threading
from typing import TYPE_CHECKING, Any, Dict, List, Optional
from uuid import UUID

from ..logger import log_event_dropped
from ..run import current_run
from ..ids import new_span_id
from .normalize import (
    evt_graph_start,
    evt_graph_end,
    evt_node_start,
    evt_node_end,
    evt_checkpoint_save,
    evt_checkpoint_load,
    evt_task_start,
    evt_task_end,
    evt_llm_start,
    evt_llm_end,
    evt_tool_start,
    evt_tool_end,
)

if TYPE_CHECKING:
    from langchain_core.outputs import LLMResult
    from ..run import ArzuleRun

# Try to import BaseCallbackHandler
try:
    from langchain_core.callbacks.base import BaseCallbackHandler

    _BASE_CLASS = BaseCallbackHandler
except ImportError:
    # LangChain/LangGraph not installed, use object as base
    _BASE_CLASS = object  # type: ignore


class ArzuleLangGraphHandler(_BASE_CLASS):
    """
    Callback handler for LangGraph that emits trace events to Arzule.

    This handler captures:
    - Graph execution start/end
    - Node execution start/end
    - Checkpoint save/load operations
    - Task execution (parallel nodes)
    - Nested LLM calls (from LangChain)
    - Nested tool calls

    Usage:
        from arzule_ingest.langgraph import instrument_langgraph
        from langgraph.graph import StateGraph

        handler = instrument_langgraph()
        
        graph = StateGraph(State)
        graph.add_node("node_a", node_a_func)
        compiled = graph.compile()
        
        # Use with invoke config
        result = compiled.invoke(input, config={"callbacks": [handler]})
    """

    def __init__(
        self,
        enable_graph: bool = True,
        enable_node: bool = True,
        enable_llm: bool = True,
        enable_tool: bool = True,
        enable_checkpoint: bool = True,
    ) -> None:
        """Initialize the handler with feature flags."""
        # Call parent init if inheriting from BaseCallbackHandler
        if _BASE_CLASS is not object:
            super().__init__()

        self.enable_graph = enable_graph
        self.enable_node = enable_node
        self.enable_llm = enable_llm
        self.enable_tool = enable_tool
        self.enable_checkpoint = enable_checkpoint

        # Track run_id -> span_id mapping for correlating start/end
        self._run_spans: Dict[str, str] = {}
        
        # Track graph execution context
        self._graph_names: Dict[str, str] = {}  # run_id -> graph_name
        self._node_names: Dict[str, str] = {}  # run_id -> node_name
        self._tool_names: Dict[str, str] = {}  # run_id -> tool_name
        
        # Cached run_id for thread fallback
        self._cached_run_id: Optional[str] = None
        self._cached_run_id_lock = threading.Lock()

    # =========================================================================
    # Run Context Management
    # =========================================================================

    def _cache_run_id(self, run_id: str) -> None:
        """Cache the run_id for thread-safe fallback lookup."""
        with self._cached_run_id_lock:
            self._cached_run_id = run_id

    def _clear_cached_run_id(self) -> None:
        """Clear the cached run_id."""
        with self._cached_run_id_lock:
            self._cached_run_id = None

    def _get_cached_run_id(self) -> Optional[str]:
        """Get the cached run_id (thread-safe)."""
        with self._cached_run_id_lock:
            return self._cached_run_id

    def _get_run_with_fallback(self, callback_name: str) -> Optional["ArzuleRun"]:
        """Get run from ContextVar, falling back to cached run_id."""
        # Try ContextVar first
        run = current_run()
        if run:
            self._cache_run_id(run.run_id)
            return run
        
        # Fallback: try global registry with cached run_id
        cached_id = self._get_cached_run_id()
        if cached_id:
            run = current_run(run_id_hint=cached_id)
            if run:
                return run
        
        # Log the drop
        if cached_id:
            log_event_dropped(
                reason="no_active_run_and_fallback_failed",
                event_class=f"langgraph.{callback_name}",
                extra={"cached_run_id": cached_id}
            )
        return None

    def _get_run_key(self, run_id: UUID) -> str:
        """Convert UUID to string key."""
        return str(run_id)

    def _start_span(self, run_id: UUID) -> str:
        """Create and track a new span for a run."""
        span_id = new_span_id()
        self._run_spans[self._get_run_key(run_id)] = span_id
        return span_id

    def _end_span(self, run_id: UUID) -> Optional[str]:
        """Get and remove the span for a run."""
        key = self._get_run_key(run_id)
        return self._run_spans.pop(key, None)

    def _get_span(self, run_id: UUID) -> Optional[str]:
        """Get the span for a run without removing it."""
        return self._run_spans.get(self._get_run_key(run_id))

    # =========================================================================
    # Chain Callbacks (used by LangGraph for graph/node execution)
    # =========================================================================

    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.
        
        In LangGraph, this is triggered for:
        - Graph execution (Pregel/CompiledStateGraph)
        - Node execution (individual node functions)
        
        Note: Only handles LangGraph-specific chains. Regular LangChain chains
        should be handled by the LangChain handler to avoid duplication.
        """
        run = self._get_run_with_fallback("on_chain_start")
        if not run:
            return

        # Determine if this is a graph or node based on the serialized info
        chain_type = self._get_chain_type(serialized)
        chain_name = self._extract_chain_name(serialized)
        
        # Skip if this is not a LangGraph-specific chain
        if chain_type == "other":
            return
        
        span_id = self._start_span(run_id)
        parent_span = self._get_span(parent_run_id) if parent_run_id else None

        if chain_type == "graph" and self.enable_graph:
            # This is a graph execution
            self._graph_names[self._get_run_key(run_id)] = chain_name
            evt = evt_graph_start(
                run=run,
                span_id=span_id,
                parent_span_id=parent_span or run.current_parent_span_id(),
                graph_name=chain_name,
                input_data=inputs,
                metadata=metadata,
            )
            run.emit(evt)
            
        elif chain_type == "node" and self.enable_node:
            # This is a node execution
            self._node_names[self._get_run_key(run_id)] = chain_name
            evt = evt_node_start(
                run=run,
                span_id=span_id,
                parent_span_id=parent_span or run.current_parent_span_id(),
                node_name=chain_name,
                input_data=inputs,
                metadata=metadata,
            )
            run.emit(evt)

    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 ends successfully."""
        run = self._get_run_with_fallback("on_chain_end")
        if not run:
            return

        span_id = self._end_span(run_id)
        parent_span = self._get_span(parent_run_id) if parent_run_id else None
        run_key = self._get_run_key(run_id)

        # Determine if this was a graph or node
        if run_key in self._graph_names and self.enable_graph:
            graph_name = self._graph_names.pop(run_key)
            evt = evt_graph_end(
                run=run,
                span_id=span_id,
                parent_span_id=parent_span or run.current_parent_span_id(),
                graph_name=graph_name,
                output_data=outputs,
            )
            run.emit(evt)
            
        elif run_key in self._node_names and self.enable_node:
            node_name = self._node_names.pop(run_key)
            evt = evt_node_end(
                run=run,
                span_id=span_id,
                parent_span_id=parent_span or run.current_parent_span_id(),
                node_name=node_name,
                output_data=outputs,
            )
            run.emit(evt)

    def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a chain encounters an error."""
        run = self._get_run_with_fallback("on_chain_error")
        if not run:
            return

        span_id = self._end_span(run_id)
        parent_span = self._get_span(parent_run_id) if parent_run_id else None
        run_key = self._get_run_key(run_id)

        # Determine if this was a graph or node
        if run_key in self._graph_names and self.enable_graph:
            graph_name = self._graph_names.pop(run_key)
            evt = evt_graph_end(
                run=run,
                span_id=span_id,
                parent_span_id=parent_span or run.current_parent_span_id(),
                graph_name=graph_name,
                output_data=None,
                error=error,
            )
            run.emit(evt)
            
        elif run_key in self._node_names and self.enable_node:
            node_name = self._node_names.pop(run_key)
            evt = evt_node_end(
                run=run,
                span_id=span_id,
                parent_span_id=parent_span or run.current_parent_span_id(),
                node_name=node_name,
                output_data=None,
                error=error,
            )
            run.emit(evt)

    # =========================================================================
    # LLM Callbacks (from nested LangChain calls)
    # =========================================================================

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[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 LLM starts running."""
        if not self.enable_llm:
            return

        run = self._get_run_with_fallback("on_llm_start")
        if not run:
            return

        span_id = self._start_span(run_id)
        parent_span = self._get_span(parent_run_id) if parent_run_id else None
        model_name = self._extract_llm_name(serialized)

        evt = evt_llm_start(
            run=run,
            span_id=span_id,
            parent_span_id=parent_span or run.current_parent_span_id(),
            model_name=model_name,
            prompts=prompts,
            metadata=metadata,
        )
        run.emit(evt)

    def on_llm_end(
        self,
        response: "LLMResult",
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when LLM finishes running."""
        if not self.enable_llm:
            return

        run = self._get_run_with_fallback("on_llm_end")
        if not run:
            return

        span_id = self._end_span(run_id)
        parent_span = self._get_span(parent_run_id) if parent_run_id else None

        # Try to extract model name from response metadata
        model_name = "unknown"
        if hasattr(response, "llm_output") and response.llm_output:
            model_name = response.llm_output.get("model_name", "unknown")

        evt = evt_llm_end(
            run=run,
            span_id=span_id,
            parent_span_id=parent_span or run.current_parent_span_id(),
            model_name=model_name,
            response=response,
        )
        run.emit(evt)

    def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when LLM encounters an error."""
        if not self.enable_llm:
            return

        run = self._get_run_with_fallback("on_llm_error")
        if not run:
            return

        span_id = self._end_span(run_id)
        parent_span = self._get_span(parent_run_id) if parent_run_id else None

        evt = evt_llm_end(
            run=run,
            span_id=span_id,
            parent_span_id=parent_span or run.current_parent_span_id(),
            model_name="unknown",
            response=None,
            error=error,
        )
        run.emit(evt)

    # =========================================================================
    # Tool Callbacks (from nested tool calls)
    # =========================================================================

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a tool starts running."""
        if not self.enable_tool:
            return

        run = self._get_run_with_fallback("on_tool_start")
        if not run:
            return

        span_id = self._start_span(run_id)
        parent_span = self._get_span(parent_run_id) if parent_run_id else None
        tool_name = serialized.get("name", "unknown")
        
        # Store tool name for retrieval in on_tool_end/on_tool_error
        self._tool_names[self._get_run_key(run_id)] = tool_name

        evt = evt_tool_start(
            run=run,
            span_id=span_id,
            parent_span_id=parent_span or run.current_parent_span_id(),
            tool_name=tool_name,
            tool_input=input_str,
        )
        run.emit(evt)

    def on_tool_end(
        self,
        output: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a tool finishes running."""
        if not self.enable_tool:
            return

        run = self._get_run_with_fallback("on_tool_end")
        if not run:
            return

        span_id = self._end_span(run_id)
        parent_span = self._get_span(parent_run_id) if parent_run_id else None
        
        # Retrieve stored tool name
        run_key = self._get_run_key(run_id)
        tool_name = self._tool_names.pop(run_key, "unknown")

        evt = evt_tool_end(
            run=run,
            span_id=span_id,
            parent_span_id=parent_span or run.current_parent_span_id(),
            tool_name=tool_name,
            output=output,
        )
        run.emit(evt)

    def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> None:
        """Called when a tool encounters an error."""
        if not self.enable_tool:
            return

        run = self._get_run_with_fallback("on_tool_error")
        if not run:
            return

        span_id = self._end_span(run_id)
        parent_span = self._get_span(parent_run_id) if parent_run_id else None
        
        # Retrieve stored tool name
        run_key = self._get_run_key(run_id)
        tool_name = self._tool_names.pop(run_key, "unknown")

        evt = evt_tool_end(
            run=run,
            span_id=span_id,
            parent_span_id=parent_span or run.current_parent_span_id(),
            tool_name=tool_name,
            output=None,
            error=error,
        )
        run.emit(evt)

    # =========================================================================
    # Helper Methods
    # =========================================================================

    def _get_chain_type(self, serialized: Dict[str, Any]) -> str:
        """Determine if this is a graph, node, or other chain based on serialized data.
        
        Returns:
            "graph" - LangGraph graph execution (Pregel/CompiledStateGraph)
            "node" - LangGraph node execution
            "other" - Regular LangChain chain (should be handled by LangChain handler)
        """
        # LangGraph graphs are typically Pregel or CompiledStateGraph
        id_list = serialized.get("id", [])
        name = serialized.get("name", "")
        repr_str = serialized.get("repr", "")
        
        # Check if this is a LangGraph graph
        if "Pregel" in str(id_list) or "CompiledStateGraph" in str(id_list):
            return "graph"
        elif "StateGraph" in str(id_list):
            return "graph"
        elif any(x in name.lower() for x in ["pregel", "graph"]):
            return "graph"
        
        # Check if this is a LangGraph node (part of a graph)
        # LangGraph nodes are typically wrapped in RunnableCallable with specific metadata
        if "RunnableCallable" in str(id_list) or "ChannelWrite" in str(id_list):
            # Check if parent is a LangGraph construct
            if any(x in str(serialized) for x in ["langgraph", "Pregel", "StateGraph"]):
                return "node"
        
        # Check if node is named explicitly
        if "node" in name.lower() and any(x in str(serialized) for x in ["langgraph", "graph"]):
            return "node"
        
        # This is a regular LangChain chain, not LangGraph-specific
        return "other"

    def _extract_chain_name(self, serialized: Dict[str, Any]) -> str:
        """Extract chain/node name from serialized data."""
        name = serialized.get("name")
        if name:
            return name

        # Try id field (list of class names)
        id_list = serialized.get("id")
        if id_list and isinstance(id_list, list) and len(id_list) > 0:
            return id_list[-1]

        return "unknown"

    def _extract_llm_name(self, serialized: Dict[str, Any]) -> str:
        """Extract LLM model name from serialized data."""
        kwargs = serialized.get("kwargs", {})
        if kwargs:
            model = kwargs.get("model_name") or kwargs.get("model") or kwargs.get("deployment_name")
            if model:
                return model

        return self._extract_chain_name(serialized)

