"""
LangGraph patches for CompiledGraph execution tracing.

Wraps graph invocation in an AGENT-kind span with orchestrator role,
capturing graph name, thread ID, and state keys. LangChain callback
handler (if LangChain integration is active) captures internal node
execution automatically.

Does NOT suppress provider instrumentation — LangGraph graphs delegate
to LangChain runnables, which handle suppression via their own patches.
"""

from __future__ import annotations

import logging
import time
from typing import Any, Callable, Dict, TypeVar

import wrapt

from risicare.integrations._base import (
    get_tracer,
    is_tracing_enabled,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
    create_framework_attributes,
    get_framework_version,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")


def _extract_graph_name(instance: Any) -> str:
    """Extract graph name from CompiledGraph instance."""
    # Try graph.name, graph.get_name(), or fallback
    name = getattr(instance, "name", None)
    if name:
        return name
    builder = getattr(instance, "builder", None)
    if builder:
        name = getattr(builder, "name", None) or getattr(builder, "__name__", None)
        if name:
            return name
    return "graph"


def _extract_state_keys(state: Any) -> list:
    """Extract state keys from graph input state."""
    if isinstance(state, dict):
        return list(state.keys())[:20]
    return []


def _wrap_graph_invoke(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for CompiledGraph.invoke (sync)."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    graph_name = _extract_graph_name(instance)
    version = get_framework_version("langgraph")
    attrs = create_framework_attributes("langgraph", version, graph_name=graph_name)
    attrs["agent.role"] = "orchestrator"
    attrs["agent.type"] = "langgraph"

    # Extract state keys from input
    state = args[0] if args else kwargs.get("input")
    if state:
        state_keys = _extract_state_keys(state)
        if state_keys:
            attrs["framework.langgraph.state_keys"] = state_keys

    # Extract thread_id from config
    config = kwargs.get("config") or {}
    configurable = config.get("configurable") or {}
    thread_id = configurable.get("thread_id")
    if thread_id:
        attrs["framework.langgraph.thread_id"] = str(thread_id)

    with tracer.start_span(
        name=f"langgraph.graph/{graph_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if should_trace_content() and isinstance(result, dict):
                result_keys = list(result.keys())[:20]
                span.set_attribute("framework.langgraph.result_keys", result_keys)

            return result
        except Exception as e:
            record_error(span, e)
            raise


async def _wrap_graph_ainvoke(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for CompiledGraph.ainvoke (async)."""
    if not is_tracing_enabled():
        return await wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    graph_name = _extract_graph_name(instance)
    version = get_framework_version("langgraph")
    attrs = create_framework_attributes("langgraph", version, graph_name=graph_name)
    attrs["agent.role"] = "orchestrator"
    attrs["agent.type"] = "langgraph"

    state = args[0] if args else kwargs.get("input")
    if state:
        state_keys = _extract_state_keys(state)
        if state_keys:
            attrs["framework.langgraph.state_keys"] = state_keys

    config = kwargs.get("config") or {}
    configurable = config.get("configurable") or {}
    thread_id = configurable.get("thread_id")
    if thread_id:
        attrs["framework.langgraph.thread_id"] = str(thread_id)

    with tracer.start_span(
        name=f"langgraph.graph/{graph_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = await wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if should_trace_content() and isinstance(result, dict):
                result_keys = list(result.keys())[:20]
                span.set_attribute("framework.langgraph.result_keys", result_keys)

            return result
        except Exception as e:
            record_error(span, e)
            raise


def _wrap_graph_stream(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for CompiledGraph.stream (sync generator)."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    graph_name = _extract_graph_name(instance)
    version = get_framework_version("langgraph")
    attrs = create_framework_attributes("langgraph", version, graph_name=graph_name)
    attrs["agent.role"] = "orchestrator"
    attrs["agent.type"] = "langgraph"
    attrs["framework.langgraph.streaming"] = True

    config = kwargs.get("config") or {}
    configurable = config.get("configurable") or {}
    thread_id = configurable.get("thread_id")
    if thread_id:
        attrs["framework.langgraph.thread_id"] = str(thread_id)

    with tracer.start_span(
        name=f"langgraph.graph/{graph_name}",
        kind=SpanKind.AGENT,
        attributes=attrs,
    ) as span:
        start_time = time.perf_counter()
        node_count = 0
        try:
            for event in wrapped(*args, **kwargs):
                node_count += 1
                yield event
        except Exception as e:
            record_error(span, e)
            raise
        finally:
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)
            span.set_attribute("framework.langgraph.streamed_events", node_count)


def patch_langgraph(module: Any) -> None:
    """Apply wrapt patches to LangGraph CompiledGraph methods."""
    try:
        wrapt.wrap_function_wrapper(
            "langgraph.graph.graph",
            "CompiledGraph.invoke",
            _wrap_graph_invoke,
        )
    except Exception as e:
        logger.debug(f"Failed to patch CompiledGraph.invoke: {e}")

    try:
        wrapt.wrap_function_wrapper(
            "langgraph.graph.graph",
            "CompiledGraph.ainvoke",
            _wrap_graph_ainvoke,
        )
    except Exception as e:
        logger.debug(f"Failed to patch CompiledGraph.ainvoke: {e}")

    try:
        wrapt.wrap_function_wrapper(
            "langgraph.graph.graph",
            "CompiledGraph.stream",
            _wrap_graph_stream,
        )
    except Exception as e:
        logger.debug(f"Failed to patch CompiledGraph.stream: {e}")
