"""
LangGraph auto-instrumentation.

Patches StateGraph.compile() to attach an Aigie callback handler via
LangGraph's standard .with_config({"callbacks": [handler]}) API.

This follows the same pattern as Langfuse and LangSmith -- callbacks are the
standard mechanism for hooking into LangGraph execution. A single handler
covers invoke, ainvoke, stream, and astream automatically.

Name extraction follows Langfuse's pattern with prioritized fallbacks:
1. inputs.metadata.workflow_name (explicit)
2. inputs.metadata.workflow_type (e.g., "deep_research" -> "deep_research_workflow")
3. Graph schema class name (e.g., ResearchState -> research_workflow)
4. Fallback to "LangGraph Workflow"
"""

import functools
import logging
import re
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

_patched_classes = set()


# ---------------------------------------------------------------------------
# Name / schema extraction (auto-specific concerns)
# ---------------------------------------------------------------------------

def _extract_workflow_name(inputs: Any, graph_schema: Any = None) -> str:
    """Extract meaningful workflow name from inputs or graph schema."""
    if isinstance(inputs, dict):
        metadata = inputs.get('metadata', {})
        if isinstance(metadata, dict):
            if metadata.get('workflow_name'):
                return metadata['workflow_name']
            if metadata.get('workflow_type'):
                wf_type = metadata['workflow_type']
                return f"{wf_type}_workflow" if not wf_type.endswith('_workflow') else wf_type
            if metadata.get('use_case'):
                return f"{metadata['use_case']}_workflow"

    if graph_schema is not None:
        schema_name = None
        if isinstance(graph_schema, type):
            schema_name = graph_schema.__name__
        elif hasattr(graph_schema, '__name__'):
            schema_name = graph_schema.__name__
        elif hasattr(graph_schema, '__class__') and graph_schema.__class__ != type:
            schema_name = graph_schema.__class__.__name__

        skip = ('dict', 'Dict', 'TypedDict', 'State', 'str', 'int', 'list', 'bool', 'float', 'NoneType')
        if schema_name and schema_name not in skip:
            snake_name = re.sub(r'(?<!^)(?=[A-Z])', '_', schema_name).lower()
            if snake_name.endswith('_state'):
                snake_name = snake_name[:-6]
            if snake_name:
                return f"{snake_name}_workflow"

    return "LangGraph Workflow"


def _extract_graph_schema(state_graph) -> Any:
    """Extract graph schema from a StateGraph instance."""
    primitives = (dict, str, int, list, bool, float)

    if hasattr(state_graph, 'schema') and state_graph.schema:
        s = state_graph.schema
        if isinstance(s, type) and s not in primitives:
            return s

    if hasattr(state_graph, '_schema') and state_graph._schema:
        s = state_graph._schema
        if isinstance(s, type) and s not in primitives:
            return s

    if hasattr(state_graph, 'schemas') and state_graph.schemas:
        for s in state_graph.schemas.keys():
            if isinstance(s, type) and s not in primitives:
                return s

    return None


# ---------------------------------------------------------------------------
# Sync batch delivery (for ephemeral event loops in invoke())
# ---------------------------------------------------------------------------

def _send_batch_sync(kytte_url: str, auth_token: str, batch_events: list) -> None:
    """Send a batch of events via synchronous HTTP POST.

    LangGraph's sync invoke() runs callbacks on an ephemeral event loop that
    closes immediately after the graph returns.  The async httpx client can't
    complete HTTP calls on a closing loop, so we use a sync httpx.Client here
    to guarantee span delivery.
    """
    if not batch_events or not kytte_url or not auth_token:
        return
    try:
        import json
        import httpx as _httpx
        payload = json.dumps({"batch": batch_events}, default=str)
        with _httpx.Client(timeout=10.0) as client:
            resp = client.post(
                f"{kytte_url}/v1/ingestion",
                content=payload.encode("utf-8"),
                headers={
                    "Content-Type": "application/json",
                    "X-API-Key": auth_token,
                },
            )
            if resp.status_code >= 400:
                logger.warning(f"Sync batch send failed ({resp.status_code}): {resp.text[:200]}")
    except Exception as e:
        logger.debug(f"Sync batch send error (non-fatal): {e}")


# ---------------------------------------------------------------------------
# Callback handler (thin wrapper around LangGraphCore)
# ---------------------------------------------------------------------------

try:
    from langchain_core.callbacks import AsyncCallbackHandler as _BaseHandler
except ImportError:
    try:
        from langchain_core.callbacks import BaseCallbackHandler as _BaseHandler
    except ImportError:
        _BaseHandler = object


class _AigieLangGraphAutoHandler(_BaseHandler):
    """LangGraph auto-instrumentation callback handler.

    Thin adapter over LangGraphCore. All span logic lives in the core;
    this class handles auto-specific concerns: compile-time patching,
    trace context management, and sync HTTP delivery.
    """

    def __init__(self, graph_schema: Any = None):
        if _BaseHandler is not object:
            super().__init__()
        self._graph_schema = graph_schema

        from ..integrations.langgraph.core import LangGraphCore
        self._core = LangGraphCore(graph_schema=graph_schema)

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

    # ------------------------------------------------------------------
    # on_chain_start / on_chain_end / on_chain_error
    # ------------------------------------------------------------------

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

        rid = str(run_id)
        pid = str(parent_run_id) if parent_run_id else None

        if parent_run_id is None:
            # ROOT: graph execution starting -> create trace
            workflow_name = _extract_workflow_name(inputs, self._graph_schema)

            from ..auto_instrument.trace import get_or_create_trace, set_current_trace, set_callback_context, clear_current_trace
            clear_current_trace()
            set_callback_context(True)

            trace = await get_or_create_trace(
                name=workflow_name,
                metadata={"type": "langgraph", "framework": "langgraph"},
            )
            if not trace:
                return
            set_current_trace(trace)

            # Direct TRACE_CREATE via HTTP to beat race condition
            if hasattr(aigie, 'client') and aigie.client and aigie._auth_token:
                try:
                    resp = await aigie.client.post(
                        f"{aigie._kytte_url}/v1/traces",
                        json={
                            "id": trace.id,
                            "name": workflow_name,
                            "metadata": {"framework": "langgraph", "type": "langgraph"},
                            "tags": [],
                        },
                        headers={"Authorization": f"Bearer {aigie._auth_token}"},
                        timeout=5.0,
                    )
                    logger.debug(f"Direct TRACE_CREATE: {resp.status_code}")
                except Exception as e:
                    logger.debug(f"Direct TRACE_CREATE failed (non-fatal): {e}")

            if aigie._buffer:
                await aigie._buffer.flush()

            self._core.start_trace(rid, trace, workflow_name, inputs)
            logger.debug(f"LangGraph auto-handler: trace {trace.id} created for '{workflow_name}'")
        else:
            # NESTED: node execution
            # Skip if parent is already a tracked node span (avoids duplicates
            # from RunnableSequence wrappers inside LangGraph nodes)
            if pid and pid in self._core._node_spans:
                return

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

            self._core.start_node(rid, pid, node_name, inputs)

    async def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)

        if parent_run_id is None and rid in self._core._runs:
            # ROOT: graph complete -> collect all buffered events + send sync
            from ..client import get_aigie
            from ..buffer import EventType
            aigie = get_aigie()

            batch_events = await self._drain_buffer(aigie)
            trace_events = self._core.end_trace(rid, outputs)
            batch_events.extend(trace_events)

            if aigie:
                _send_batch_sync(aigie._kytte_url, aigie._auth_token, batch_events)

            from ..auto_instrument.trace import clear_current_trace, set_callback_context
            clear_current_trace()
            set_callback_context(False)

        elif rid in self._core._node_spans:
            # NODE: complete -> buffer the span
            span = self._core.end_node(rid, outputs)
            if span:
                from ..client import get_aigie
                from ..buffer import EventType
                aigie = get_aigie()
                if aigie and aigie._buffer:
                    await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def on_chain_error(
        self,
        error: BaseException,
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)

        if parent_run_id is None and rid in self._core._runs:
            from ..client import get_aigie
            aigie = get_aigie()

            batch_events = await self._drain_buffer(aigie)
            trace_events = self._core.error_trace(rid, error)
            batch_events.extend(trace_events)

            if aigie:
                _send_batch_sync(aigie._kytte_url, aigie._auth_token, batch_events)

            from ..auto_instrument.trace import clear_current_trace, set_callback_context
            clear_current_trace()
            set_callback_context(False)

        elif rid in self._core._node_spans:
            span = self._core.error_node(rid, error)
            if span:
                from ..client import get_aigie
                from ..buffer import EventType
                aigie = get_aigie()
                if aigie and aigie._buffer:
                    await aigie._buffer.add(EventType.SPAN_CREATE, span)

    # ------------------------------------------------------------------
    # on_chat_model_start / on_llm_start / on_llm_end / on_llm_error
    # ------------------------------------------------------------------

    async def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[Any],
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        # Extract structured messages (role/content) instead of stringified repr
        structured = []
        for msg_list in messages:
            items = msg_list if isinstance(msg_list, list) else [msg_list]
            for m in items:
                role = getattr(m, "type", "unknown")
                # Map LangChain types to standard roles
                if role == "human":
                    role = "user"
                elif role == "ai":
                    role = "assistant"
                content = getattr(m, "content", str(m))
                entry: Dict[str, Any] = {"role": role, "content": str(content)[:1000]}
                # Include tool calls if present
                tool_calls = getattr(m, "tool_calls", None)
                if tool_calls:
                    entry["tool_calls"] = [
                        {"name": tc.get("name", ""), "args": tc.get("args", {})}
                        for tc in tool_calls[:10]
                    ]
                structured.append(entry)
        await self.on_llm_start(serialized, structured, run_id=run_id, parent_run_id=parent_run_id, **kwargs)

    async def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: Any,
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        from ..client import get_aigie
        aigie = 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")
        )
        # Strip provider prefixes (e.g. models/gemini-2.5-flash -> gemini-2.5-flash)
        if model_name.startswith("models/"):
            model_name = model_name[7:]

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

    async def on_llm_end(
        self,
        response: Any,
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)
        span = self._core.end_llm(rid, response)
        if span:
            from ..client import get_aigie
            from ..buffer import EventType
            aigie = get_aigie()
            if aigie and aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_CREATE, span)
            # Run FastDetector for signal hints (passed to background judge)
            _signal_hints = self._run_fast_detect(aigie, span)
            self._fire_judge(aigie, span, signal_hints=_signal_hints)

    async def on_llm_error(
        self,
        error: BaseException,
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)
        span = self._core.error_llm(rid, error)
        if span:
            from ..client import get_aigie
            from ..buffer import EventType
            aigie = get_aigie()
            if aigie and aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_CREATE, span)
            self._fire_judge(aigie, span, error=str(error))

    @staticmethod
    def _run_fast_detect(aigie: Any, span: Dict[str, Any]) -> List[str]:
        """Run FastDetector for signal hints to guide judge selection."""
        detector = getattr(aigie, '_fast_detector', None) if aigie else None
        if not detector:
            return []
        try:
            output = span.get("output", {})
            content = output.get("content", "") if isinstance(output, dict) else str(output)
            meta = span.get("metadata", {})
            token_usage = meta.get("token_usage", {})
            result = detector.detect(
                trace_id=span.get("trace_id", ""),
                span_id=span.get("id", ""),
                output_content=content,
                model=meta.get("model"),
                input_tokens=token_usage.get("input_tokens", 0),
                output_tokens=token_usage.get("output_tokens", 0),
                cost=meta.get("cost", 0.0),
            )
            return result.signal_hints if result else []
        except Exception:
            return []

    @staticmethod
    def _fire_judge(
        aigie: Any,
        span: Dict[str, Any],
        error: Optional[str] = None,
        signal_hints: Optional[List[str]] = None,
    ) -> None:
        """Fire async judge evaluation for a completed LLM span.

        Always schedules on the persistent bg loop (not the LangGraph callback
        loop) to avoid cross-loop issues with httpx clients created during init.
        """
        import asyncio
        runner = getattr(aigie, '_async_judge_runner', None) if aigie else None
        if not runner:
            return
        bg_loop = getattr(aigie, '_bg_loop', None)
        if not bg_loop or not bg_loop.is_running():
            return
        try:
            output = span.get("output", {})
            output_content = output.get("content", "") if isinstance(output, dict) else str(output)
            input_data = span.get("input", {})
            input_messages = input_data.get("messages", []) if isinstance(input_data, dict) else []
            coro = runner._run_evaluation(
                trace_id=span.get("trace_id", ""),
                span_id=span.get("id", ""),
                span_type=span.get("type", "llm"),
                output_content=output_content,
                input_messages=input_messages,
                metadata=span.get("metadata", {}),
                error=error,
                signal_hints=signal_hints,
            )
            asyncio.run_coroutine_threadsafe(coro, bg_loop)
        except Exception:
            pass  # Never block the hot path

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

    # ------------------------------------------------------------------
    # on_tool_start / on_tool_end / on_tool_error
    # ------------------------------------------------------------------

    async def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        from ..client import get_aigie
        aigie = get_aigie()
        if not aigie or not aigie._initialized:
            return

        tool_name = serialized.get("name") or kwargs.get("name", "tool")
        rid = str(run_id)
        pid = str(parent_run_id) if parent_run_id else None
        self._core.start_tool(rid, pid, tool_name, input_str)

    async def on_tool_end(
        self,
        output: str,
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)
        span = self._core.end_tool(rid, output)
        if span:
            from ..client import get_aigie
            from ..buffer import EventType
            aigie = get_aigie()
            if aigie and aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def on_tool_error(
        self,
        error: BaseException,
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)
        span = self._core.error_tool(rid, error)
        if span:
            from ..client import get_aigie
            from ..buffer import EventType
            aigie = get_aigie()
            if aigie and aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_CREATE, span)

    # ------------------------------------------------------------------
    # on_retriever_start / on_retriever_end / on_retriever_error (NEW)
    # ------------------------------------------------------------------

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

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

    async def on_retriever_end(
        self,
        documents: Any,
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)
        span = self._core.end_retriever(rid, documents)
        if span:
            from ..client import get_aigie
            from ..buffer import EventType
            aigie = get_aigie()
            if aigie and aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_CREATE, span)

    async def on_retriever_error(
        self,
        error: BaseException,
        *,
        run_id: Any,
        parent_run_id: Any = None,
        **kwargs: Any,
    ) -> None:
        rid = str(run_id)
        span = self._core.error_retriever(rid, error)
        if span:
            from ..client import get_aigie
            from ..buffer import EventType
            aigie = get_aigie()
            if aigie and aigie._buffer:
                await aigie._buffer.add(EventType.SPAN_CREATE, span)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    async def _drain_buffer(self, aigie) -> list:
        """Drain buffered events and convert to batch format."""
        if not aigie or not aigie._buffer:
            return []

        from datetime import datetime, timezone
        pending_events = []
        async with aigie._buffer._lock:
            pending_events = list(aigie._buffer._buffer)
            aigie._buffer._buffer.clear()

        batch_events = []
        evt_type_map = {
            "span_create": "span-create", "span_update": "span-update",
            "trace_create": "trace-create", "trace_update": "trace-update",
        }
        now_iso = datetime.now(timezone.utc).isoformat()
        for evt in pending_events:
            etype = evt_type_map.get(evt.event_type.value, evt.event_type.value)
            eid = evt.payload.get("id") or evt.payload.get("trace_id", "")
            body = evt.payload.copy()
            ts = body.get("start_time") or body.get("timestamp") or now_iso
            batch_events.append({"type": etype, "id": str(eid), "timestamp": str(ts), "body": body})

        return batch_events


# ---------------------------------------------------------------------------
# Patching
# ---------------------------------------------------------------------------

def patch_langgraph() -> None:
    """Patch LangGraph StateGraph.compile() for auto-instrumentation."""
    _patch_state_graph()


def _patch_state_graph() -> None:
    """Patch StateGraph.compile() to attach Aigie callbacks via .with_config()."""
    try:
        from langgraph.graph import StateGraph

        if StateGraph in _patched_classes:
            return

        original_compile = StateGraph.compile

        @functools.wraps(original_compile)
        def traced_compile(self, **kwargs):
            app = original_compile(self, **kwargs)

            graph_schema = _extract_graph_schema(self)
            handler = _AigieLangGraphAutoHandler(graph_schema=graph_schema)
            app = app.with_config({"callbacks": [handler]})

            logger.info("LangGraph auto-instrumentation: callbacks attached via with_config()")
            return app

        StateGraph.compile = traced_compile
        _patched_classes.add(StateGraph)

        logger.debug("Patched StateGraph.compile for auto-instrumentation")

    except ImportError:
        pass  # LangGraph not installed
    except Exception as e:
        logger.warning(f"Failed to patch StateGraph: {e}")
