"""
LangGraph integration for Aigie SDK.

Provides the manual-path callback handler and utility wrappers/decorators
for LangGraph applications. Delegates all span logic to the shared
LangGraphCore so both auto and manual paths produce identical telemetry.
"""

import logging
import uuid
from datetime import datetime, timezone
from functools import wraps
from typing import Any, Callable, Dict, List, TypeVar

from .buffer import EventType
from .context_manager import get_current_trace_context

try:
    from langchain_core.callbacks import BaseCallbackHandler

    HAS_BASE_CALLBACK = True
except ImportError:
    HAS_BASE_CALLBACK = False
    BaseCallbackHandler = object  # Fallback

logger = logging.getLogger(__name__)
T = TypeVar("T", bound=Callable[..., Any])


class LangGraphHandler(BaseCallbackHandler if HAS_BASE_CALLBACK else object):
    """LangGraph callback handler for Aigie.

    Automatically traces LangGraph workflow executions. Can be used
    standalone or together with an existing trace context.

    Example::

        handler = LangGraphHandler(trace_name='my-graph')
        result = await app.ainvoke(
            {"input": "..."},
            config={"callbacks": [handler]},
        )
    """

    def __init__(
        self,
        trace_name: str | None = None,
        metadata: Dict[str, Any] | None = None,
        tags: List[str] | None = None,
    ):
        if HAS_BASE_CALLBACK:
            super().__init__()

        self.trace_name = trace_name
        self.metadata = metadata or {}
        self.tags = tags or []
        self.trace_id: str | None = None
        self._aigie = None
        self._trace_context: Any | None = None
        self._langgraph_span_id: str | None = None

        # BaseCallbackHandler compat
        self.run_inline = False
        self.raise_error = False

        # Core engine (lazy-init to avoid import cycle)
        self._core = None

    def _get_core(self):
        if self._core is None:
            from .integrations.langgraph.core import LangGraphCore

            self._core = LangGraphCore()
        return self._core

    def _get_aigie(self):
        if self._aigie is None:
            from .client import get_aigie

            self._aigie = get_aigie()
        return self._aigie

    # ------------------------------------------------------------------
    # High-level API (handle_graph_start / handle_node_start / etc.)
    # ------------------------------------------------------------------

    async def handle_graph_start(self, graph_name: str, input: Any) -> None:
        """Called when graph execution starts."""
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        if self._trace_context:
            from .auto_instrument.trace import set_current_trace

            set_current_trace(self._trace_context)

        self._resolved_trace_name = self.trace_name or graph_name

        if not self.trace_id:
            if self._trace_context and hasattr(self._trace_context, "id"):
                self.trace_id = str(self._trace_context.id)
            else:
                self.trace_id = str(uuid.uuid4())

        # Create a lightweight trace-like object for the core
        trace_obj = type("_Trace", (), {"id": self.trace_id})()

        core = self._get_core()
        root_run_id = f"manual-{self.trace_id}"
        core.start_trace(
            root_run_id,
            trace_obj,
            self._resolved_trace_name,
            input,
        )
        self._root_run_id = root_run_id

        # Create trace event if no external trace context
        if not self._trace_context:
            trace_data = {
                "id": self.trace_id,
                "name": self._resolved_trace_name,
                "type": "chain",
                "input": input,
                "status": "pending",
                "tags": [*self.tags, "langgraph"],
                "metadata": {
                    **self.metadata,
                    "graphName": graph_name,
                    "framework": "langgraph",
                },
                "start_time": datetime.now(timezone.utc).isoformat(),
                "created_at": datetime.now(timezone.utc).isoformat(),
            }
            if aigie._buffer:
                await aigie._buffer.add(EventType.TRACE_CREATE, trace_data)

    async def handle_node_start(self, node_name: str, input: Any) -> None:
        """Called when a node starts executing."""
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return

        if self._trace_context:
            from .auto_instrument.trace import set_current_trace

            set_current_trace(self._trace_context)

        core = self._get_core()
        node_run_id = f"node-{uuid.uuid4()}"
        root_run_id = getattr(self, "_root_run_id", f"manual-{self.trace_id}")
        core.start_node(node_run_id, root_run_id, node_name, input)

        # Store mapping for handle_node_end
        if not hasattr(self, "_node_run_ids"):
            self._node_run_ids: Dict[str, str] = {}
        self._node_run_ids[node_name] = node_run_id

    async def handle_node_end(self, node_name: str, output: Any) -> None:
        """Called when a node finishes executing."""
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        node_run_ids = getattr(self, "_node_run_ids", {})
        node_run_id = node_run_ids.pop(node_name, None)
        if not node_run_id:
            return

        core = self._get_core()
        span = core.end_node(node_run_id, output)
        if span and aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def handle_node_error(self, node_name: str, error: Exception) -> None:
        """Called when a node encounters an error."""
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        node_run_ids = getattr(self, "_node_run_ids", {})
        node_run_id = node_run_ids.pop(node_name, None)
        if not node_run_id:
            return

        core = self._get_core()
        span = core.error_node(node_run_id, error)
        if span and aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def handle_graph_end(self, output: Any) -> None:
        """Called when graph execution finishes."""
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return

        core = self._get_core()
        root_run_id = getattr(self, "_root_run_id", f"manual-{self.trace_id}")
        trace_events = core.end_trace(root_run_id, output)

        # Send workflow span + trace update via buffer
        if aigie._buffer:
            for evt in trace_events:
                evt_type = evt.get("type", "")
                if evt_type == "span-create":
                    await aigie._buffer.add(EventType.SPAN_CREATE, evt["body"])
                elif evt_type == "trace-update":
                    await aigie._buffer.add(EventType.TRACE_UPDATE, evt["body"])

            # Flush, then re-send TRACE_CREATE to restore name
            await aigie._buffer.flush()
            resolved_name = (
                getattr(self, "_resolved_trace_name", None) or self.trace_name or "LangGraph"
            )
            await aigie._buffer.add(
                EventType.TRACE_CREATE, {"id": self.trace_id, "name": resolved_name}
            )
            await aigie._buffer.flush()

    async def handle_graph_error(self, error: Exception) -> None:
        """Called when graph execution encounters an error."""
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return

        core = self._get_core()
        root_run_id = getattr(self, "_root_run_id", f"manual-{self.trace_id}")
        trace_events = core.error_trace(root_run_id, error)

        if aigie._buffer:
            for evt in trace_events:
                evt_type = evt.get("type", "")
                if evt_type == "span-create":
                    await aigie._buffer.add(EventType.SPAN_CREATE, evt["body"])
                elif evt_type == "trace-update":
                    await aigie._buffer.add(EventType.TRACE_UPDATE, evt["body"])

    # ------------------------------------------------------------------
    # BaseCallbackHandler protocol (on_chain_*, on_llm_*, on_tool_*)
    # These delegate to the core via synthetic run IDs.
    # ------------------------------------------------------------------

    async def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: str,
        parent_run_id: str | None = None,
        tags: List[str] | None = None,
        metadata: Dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        node_name = None
        if serialized:
            node_name = serialized.get("name") or serialized.get("id")
        if not node_name and metadata:
            node_name = metadata.get("langgraph_node")
        if not node_name:
            node_name = kwargs.get("name", "Chain")

        if node_name in ("RunnableSequence", "CompiledStateGraph", "Pregel"):
            return

        if not self.trace_id:
            if self._trace_context and hasattr(self._trace_context, "id"):
                self.trace_id = str(self._trace_context.id)
            else:
                return

        core = self._get_core()
        rid = str(run_id)
        pid = str(parent_run_id) if parent_run_id else getattr(self, "_root_run_id", None)
        core.start_node(rid, pid, node_name, inputs)

    async def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: str,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        core = self._get_core()
        span = core.end_node(str(run_id), outputs)
        if span and aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def on_chain_error(
        self,
        error: Exception,
        *,
        run_id: str,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        core = self._get_core()
        span = core.error_node(str(run_id), error)
        if span and aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: str,
        parent_run_id: str | None = None,
        tags: List[str] | None = None,
        metadata: Dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        model_name = (
            kwargs.get("invocation_params", {}).get("model_name")
            or kwargs.get("invocation_params", {}).get("model")
            or (serialized.get("name", "LLM") if serialized else "LLM")
        )

        core = self._get_core()
        rid = str(run_id)
        pid = str(parent_run_id) if parent_run_id else None
        core.start_llm(rid, pid, model_name, prompts)

    async def on_llm_end(
        self,
        response: Any,
        *,
        run_id: str,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        core = self._get_core()
        span = core.end_llm(str(run_id), response)
        if span and aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def on_llm_new_token(
        self,
        token: str,
        *,
        run_id: str,
        **kwargs: Any,
    ) -> None:
        core = self._get_core()
        core.llm_new_token(str(run_id), token)

    async def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: str,
        parent_run_id: str | None = None,
        tags: List[str] | None = None,
        metadata: Dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized or not self.trace_id:
            return

        tool_name = serialized.get("name", "Tool") if serialized else "Tool"
        core = self._get_core()
        rid = str(run_id)
        pid = str(parent_run_id) if parent_run_id else None
        core.start_tool(rid, pid, tool_name, input_str)

        # Pre-tool interception signal (observe only, never blocks)
        try:
            intercept = await aigie.intercept_before_tool(
                tool_name=tool_name,
                tool_args={"input": input_str},
                trace_id=self.trace_id,
            )
            if intercept.get("decision") != "allow":
                logger.info(f"[AIGIE] Pre-tool signal for {tool_name}: {intercept.get('decision')}")
        except Exception:
            pass

    async def on_tool_end(
        self,
        output: str,
        *,
        run_id: str,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        core = self._get_core()
        span = core.end_tool(str(run_id), output)
        if span and aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_CREATE, span)

        # Post-tool interception: report errors to backend for pattern learning
        if span and span.get("is_error"):
            try:
                await aigie.intercept_after_tool(
                    tool_name=span.get("name", "Tool"),
                    result=output,
                    error=span.get("error"),
                    error_type=span.get("error_type"),
                    trace_id=self.trace_id,
                    duration_ms=span.get("metadata", {}).get("duration_ms", 0),
                )
            except Exception:
                pass

    async def on_retriever_start(
        self,
        serialized: Dict[str, Any],
        query: str,
        *,
        run_id: str,
        parent_run_id: str | None = None,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        core = self._get_core()
        rid = str(run_id)
        pid = str(parent_run_id) if parent_run_id else None
        core.start_retriever(rid, pid, query, serialized)

    async def on_retriever_end(
        self,
        documents: Any,
        *,
        run_id: str,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        core = self._get_core()
        span = core.end_retriever(str(run_id), documents)
        if span and aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def on_retriever_error(
        self,
        error: BaseException,
        *,
        run_id: str,
        **kwargs: Any,
    ) -> None:
        aigie = self._get_aigie()
        if not aigie or not aigie._initialized:
            return

        core = self._get_core()
        span = core.error_retriever(str(run_id), error)
        if span and aigie._buffer:
            await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def flush(self) -> None:
        """Clean up resources."""
        self.trace_id = None

    @staticmethod
    def _is_error_result(output: Any) -> bool:
        """Check if a tool output represents an error."""
        if output is None:
            return False
        output_str = str(output).lower()
        error_indicators = ("error:", "exception:", "traceback", "failed:", "status: error")
        return any(indicator in output_str for indicator in error_indicators)


# ---------------------------------------------------------------------------
# Utility wrappers / decorators
# ---------------------------------------------------------------------------


def wrap_langgraph(
    app: Any,
    trace_name: str | None = None,
    metadata: Dict[str, Any] | None = None,
    tags: List[str] | None = None,
) -> Any:
    """Wrap LangGraph application for automatic tracing.

    Example::

        traced_app = wrap_langgraph(app, trace_name='my-graph')
        result = await traced_app.ainvoke(input)
    """
    from .client import get_aigie

    aigie = get_aigie()
    if not aigie or not aigie._initialized:
        return app

    original_ainvoke = getattr(app, "ainvoke", None)
    if original_ainvoke:

        async def traced_ainvoke(input: Any, config: Dict[str, Any] | None = None, **kwargs):
            return await aigie.trace(
                trace_name or "langgraph",
                lambda: original_ainvoke(input, config=config, **kwargs),
                type="chain",
                input=input,
                tags=[*(tags or []), "langgraph"],
                metadata={**(metadata or {}), "framework": "langgraph"},
            )

        app.ainvoke = traced_ainvoke

    original_invoke = getattr(app, "invoke", None)
    if original_invoke:

        def traced_invoke(input: Any, config: Dict[str, Any] | None = None, **kwargs):
            return original_invoke(input, config=config, **kwargs)

        app.invoke = traced_invoke

    original_astream = getattr(app, "astream", None)
    if original_astream:

        async def traced_astream(input: Any, config: Dict[str, Any] | None = None, **kwargs):
            async for chunk in original_astream(input, config=config, **kwargs):
                yield chunk

        app.astream = traced_astream

    return app


def trace_langgraph_node(
    node_name: str,
    metadata: Dict[str, Any] | None = None,
    tags: List[str] | None = None,
) -> Callable[[T], T]:
    """Decorator to trace a LangGraph node function.

    Example::

        @trace_langgraph_node('my-node')
        async def my_node(state):
            return {'result': 'done'}
    """
    from .client import get_aigie

    def decorator(fn: T) -> T:
        @wraps(fn)
        async def async_wrapper(*args, **kwargs):
            aigie = get_aigie()
            if not aigie or not aigie._initialized:
                return await fn(*args, **kwargs)
            return await aigie.span(
                f"node:{node_name}",
                lambda: fn(*args, **kwargs),
                type="chain",
                input=args[0] if args else None,
                tags=[*(tags or []), "langgraph", "node"],
                metadata={**(metadata or {}), "nodeName": node_name, "nodeType": "langgraph_node"},
            )

        @wraps(fn)
        def sync_wrapper(*args, **kwargs):
            return fn(*args, **kwargs)

        import inspect

        if inspect.iscoroutinefunction(fn):
            return async_wrapper  # type: ignore
        else:
            return sync_wrapper  # type: ignore

    return decorator


def trace_langgraph_edge(
    edge_name: str,
    metadata: Dict[str, Any] | None = None,
) -> Callable[[T], T]:
    """Decorator to trace a LangGraph edge condition.

    Example::

        @trace_langgraph_edge('my-condition')
        def should_continue(state):
            return 'next_node' if state['count'] < 10 else 'end'
    """
    from .client import get_aigie

    def decorator(fn: T) -> T:
        @wraps(fn)
        def wrapper(*args, **kwargs):
            aigie = get_aigie()
            if not aigie or not aigie._initialized:
                return fn(*args, **kwargs)

            result = fn(*args, **kwargs)

            context = get_current_trace_context()
            if context and context.get("spanId"):
                import asyncio

                try:
                    loop = asyncio.get_event_loop()
                    if loop.is_running():
                        pass
                    else:
                        loop.run_until_complete(
                            aigie._update_span(
                                context["spanId"],
                                {
                                    "metadata": {
                                        "edgeName": edge_name,
                                        "edgeResult": result,
                                        **(metadata or {}),
                                    },
                                },
                            )
                        )
                except RuntimeError:
                    pass

            return result

        return wrapper  # type: ignore

    return decorator


def create_langgraph_handler(
    trace_name: str | None = None,
    metadata: Dict[str, Any] | None = None,
    tags: List[str] | None = None,
) -> LangGraphHandler:
    """Create LangGraph handler."""
    return LangGraphHandler(
        trace_name=trace_name,
        metadata=metadata,
        tags=tags,
    )
