"""
LangGraph auto-instrumentation.

Automatically patches LangGraph workflows to create traces and inject handlers.

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

logger = logging.getLogger(__name__)

_patched_classes = set()


def _safe_add_callback(config: dict, handler) -> None:
    """Safely add a callback handler to config, regardless of callback type.

    Handles plain lists, AsyncCallbackManager, CallbackManager, and other types.
    """
    existing = config.get("callbacks")
    if existing is None:
        config["callbacks"] = [handler]
    elif isinstance(existing, list):
        existing.append(handler)
    else:
        # AsyncCallbackManager or CallbackManager — extract handlers and convert to list
        handlers = list(getattr(existing, "handlers", []))
        handlers += list(getattr(existing, "inheritable_handlers", []))
        config["callbacks"] = handlers + [handler]


def _extract_workflow_name(inputs: Any, graph_schema: Any = None) -> str:
    """Extract meaningful workflow name from inputs or graph schema.

    Follows Langfuse's pattern with prioritized fallbacks.
    """
    logger.debug(
        f"_extract_workflow_name called with inputs type: {type(inputs)}, graph_schema: {graph_schema}"
    )

    # Check inputs metadata first
    if isinstance(inputs, dict):
        metadata = inputs.get("metadata", {})
        logger.debug(f"Extracting workflow name - inputs.metadata: {metadata}")
        if isinstance(metadata, dict):
            # 1. Explicit workflow name
            if metadata.get("workflow_name"):
                logger.debug(f"Using workflow_name from metadata: {metadata['workflow_name']}")
                return metadata["workflow_name"]
            # 2. Workflow type -> workflow name
            if metadata.get("workflow_type"):
                wf_type = metadata["workflow_type"]
                result = f"{wf_type}_workflow" if not wf_type.endswith("_workflow") else wf_type
                logger.debug(f"Using workflow_type from metadata: {wf_type} -> {result}")
                return result
            # 3. Use case -> workflow name
            if metadata.get("use_case"):
                result = f"{metadata['use_case']}_workflow"
                logger.debug(f"Using use_case from metadata: {result}")
                return result
        else:
            logger.debug(f"metadata is not a dict: {type(metadata)}")
    else:
        logger.debug(f"inputs is not a dict: {type(inputs)}")

    # 4. Extract from graph schema class name
    if graph_schema is not None:
        logger.debug(f"Extracting from graph_schema: {graph_schema}, type: {type(graph_schema)}")
        schema_name = None

        # Only use __name__ if it's actually a class/type
        if isinstance(graph_schema, type):
            schema_name = graph_schema.__name__
            logger.debug(f"Got schema name from type.__name__: {schema_name}")
        elif hasattr(graph_schema, "__name__"):
            schema_name = graph_schema.__name__
            logger.debug(f"Got schema name from __name__: {schema_name}")
        elif hasattr(graph_schema, "__class__") and graph_schema.__class__ != type:
            schema_name = graph_schema.__class__.__name__
            logger.debug(f"Got schema name from __class__.__name__: {schema_name}")

        # Filter out built-in types and common non-meaningful names
        if schema_name and schema_name not in (
            "dict",
            "Dict",
            "TypedDict",
            "State",
            "str",
            "int",
            "list",
            "bool",
            "float",
            "NoneType",
        ):
            # Convert PascalCase to snake_case
            snake_name = re.sub(r"(?<!^)(?=[A-Z])", "_", schema_name).lower()
            # Remove 'state' suffix if present
            if snake_name.endswith("_state"):
                snake_name = snake_name[:-6]
            if snake_name:
                result = f"{snake_name}_workflow"
                logger.debug(f"Using schema-derived name: {result}")
                return result
        else:
            logger.debug(f"Schema name filtered out: {schema_name}")

    # 5. Fallback
    logger.debug("Using fallback name: LangGraph Workflow")
    return "LangGraph Workflow"


def patch_langgraph() -> None:
    """Patch LangGraph classes for auto-instrumentation."""
    _patch_state_graph()
    _patch_callback_manager()
    # NOTE: _patch_compiled_graph() removed — causes infinite recursion when
    # combined with compile-time patching and langchain Runnable.ainvoke patch.


def _patch_state_graph() -> None:
    """Patch StateGraph.compile() to return auto-instrumented app."""
    try:
        from langgraph.graph import StateGraph

        if StateGraph in _patched_classes:
            return
        if getattr(StateGraph.compile, "_aigie_patched", False):
            return

        original_compile = StateGraph.compile

        @functools.wraps(original_compile)
        def traced_compile(self, **kwargs):
            """Traced version of compile."""
            app = original_compile(self, **kwargs)

            # Store reference to graph schema for name extraction
            # The schema should be the TypedDict/class used to define the state
            graph_schema = None

            # Method 1: Direct schema attribute (preferred - set by StateGraph.__init__)
            if hasattr(self, "schema") and self.schema:
                schema = self.schema
                # Make sure it's actually a class/type, not a string or primitive
                if isinstance(schema, type) and schema not in (dict, str, int, list, bool, float):
                    graph_schema = schema
                    logger.debug(f"Got graph_schema from self.schema: {graph_schema}")

            # Method 2: Try _schema (internal attribute)
            if not graph_schema and hasattr(self, "_schema") and self._schema:
                schema = self._schema
                if isinstance(schema, type) and schema not in (dict, str, int, list, bool, float):
                    graph_schema = schema
                    logger.debug(f"Got graph_schema from self._schema: {graph_schema}")

            # Method 3: Check 'schemas' attribute (LangGraph >= 0.2.x stores state class as key)
            if not graph_schema and hasattr(self, "schemas") and self.schemas:
                # self.schemas is a dict with the state class as keys
                for schema_class in self.schemas.keys():
                    if isinstance(schema_class, type) and schema_class not in (
                        dict,
                        str,
                        int,
                        list,
                        bool,
                        float,
                    ):
                        graph_schema = schema_class
                        logger.debug(f"Got graph_schema from self.schemas: {graph_schema}")
                        break

            # Method 4: Check channels for the state type (fallback for older versions)
            if not graph_schema and hasattr(self, "channels") and self.channels:
                logger.debug(f"channels.keys(): {list(self.channels.keys())}")

            # Patch the compiled app's invoke methods
            if hasattr(app, "ainvoke") and not getattr(app.ainvoke, "_aigie_patched", False):
                original_ainvoke = app.ainvoke

                async def traced_ainvoke(
                    input: Any, config: Dict[str, Any] | None = None, **kwargs
                ):
                    from ..auto_instrument.trace import (
                        get_or_create_trace,
                        set_active_callback_handler,
                        set_callback_context,
                        set_current_trace,
                    )
                    from ..callback import AigieCallbackHandler
                    from ..client import get_aigie

                    aigie = get_aigie()

                    # Ensure aigie is initialized
                    if aigie and not aigie._initialized:
                        try:
                            await aigie.initialize()
                        except Exception as e:
                            logger.debug(f"Failed to initialize aigie: {e}")

                    if not aigie or not aigie._initialized:
                        return await original_ainvoke(input, config=config, **kwargs)

                    # Extract workflow name from inputs/schema
                    workflow_name = _extract_workflow_name(input, graph_schema)

                    # Create trace
                    trace = await get_or_create_trace(
                        name=workflow_name,
                        metadata={
                            "type": "langgraph",
                            "framework": "langgraph",
                            "inputs": input if isinstance(input, dict) else {},
                        },
                    )
                    set_current_trace(trace)

                    # ONE handler — the universal pattern
                    handler = AigieCallbackHandler(aigie=aigie, trace=trace)
                    handler._langgraph_mode = True
                    handler._workflow_name = workflow_name

                    if config is None:
                        config = {}
                    _safe_add_callback(config, handler)

                    # Set run_name in config metadata
                    if "metadata" not in config:
                        config["metadata"] = {}
                    config["metadata"]["run_name"] = workflow_name

                    # Suppress LLM auto-instrumentation (callbacks handle it)
                    # and store the handler so BaseCallbackManager.__init__ patch
                    # can auto-inject it into new callback managers (for LLM calls
                    # inside nodes that don't pass config).
                    set_callback_context(True)
                    set_active_callback_handler(handler)
                    try:
                        return await original_ainvoke(input, config=config, **kwargs)
                    finally:
                        set_active_callback_handler(None)
                        set_callback_context(False)

                traced_ainvoke._aigie_patched = True
                app.ainvoke = traced_ainvoke

            if hasattr(app, "invoke") and not getattr(app.invoke, "_aigie_patched", False):
                original_invoke = app.invoke

                def traced_invoke(input: Any, config: Dict[str, Any] | None = None, **kwargs):
                    """Traced version of invoke."""
                    from ..auto_instrument.trace import (
                        get_or_create_trace_sync,
                        set_active_callback_handler,
                        set_callback_context,
                        set_current_trace,
                    )
                    from ..callback import AigieCallbackHandler
                    from ..client import get_aigie

                    aigie = get_aigie()
                    if not aigie or not aigie._initialized:
                        return original_invoke(input, config=config, **kwargs)

                    workflow_name = _extract_workflow_name(input, graph_schema)
                    trace = get_or_create_trace_sync(
                        name=workflow_name,
                        metadata={
                            "type": "langgraph",
                            "framework": "langgraph",
                            "inputs": input if isinstance(input, dict) else {},
                        },
                    )

                    handler = None
                    if trace:
                        set_current_trace(trace)
                        handler = AigieCallbackHandler(aigie=aigie, trace=trace)
                        handler._langgraph_mode = True
                        handler._workflow_name = workflow_name

                        if config is None:
                            config = {}
                        _safe_add_callback(config, handler)

                        if "metadata" not in config:
                            config["metadata"] = {}
                        config["metadata"]["run_name"] = workflow_name

                    set_callback_context(True)
                    set_active_callback_handler(handler)
                    try:
                        return original_invoke(input, config=config, **kwargs)
                    finally:
                        set_active_callback_handler(None)
                        set_callback_context(False)

                traced_invoke._aigie_patched = True
                app.invoke = traced_invoke

            if hasattr(app, "astream") and not getattr(app.astream, "_aigie_patched", False):
                original_astream = app.astream

                async def traced_astream(
                    input: Any, config: Dict[str, Any] | None = None, **kwargs
                ):
                    from ..auto_instrument.trace import (
                        get_or_create_trace,
                        set_active_callback_handler,
                        set_callback_context,
                        set_current_trace,
                    )
                    from ..callback import AigieCallbackHandler
                    from ..client import get_aigie

                    aigie = get_aigie()

                    if aigie and not aigie._initialized:
                        try:
                            await aigie.initialize()
                        except Exception as e:
                            logger.debug(f"Failed to initialize aigie: {e}")

                    if not aigie or not aigie._initialized:
                        async for chunk in original_astream(input, config=config, **kwargs):
                            yield chunk
                        return

                    workflow_name = _extract_workflow_name(input, graph_schema)

                    trace = await get_or_create_trace(
                        name=workflow_name,
                        metadata={
                            "type": "langgraph",
                            "framework": "langgraph",
                            "inputs": input if isinstance(input, dict) else {},
                        },
                    )
                    set_current_trace(trace)

                    handler = AigieCallbackHandler(aigie=aigie, trace=trace)
                    handler._langgraph_mode = True
                    handler._workflow_name = workflow_name

                    if config is None:
                        config = {}
                    _safe_add_callback(config, handler)

                    if "metadata" not in config:
                        config["metadata"] = {}
                    config["metadata"]["run_name"] = workflow_name

                    set_callback_context(True)
                    set_active_callback_handler(handler)
                    try:
                        async for chunk in original_astream(input, config=config, **kwargs):
                            yield chunk
                    finally:
                        set_active_callback_handler(None)
                        set_callback_context(False)

                traced_astream._aigie_patched = True
                app.astream = traced_astream

            if hasattr(app, "stream") and not getattr(app.stream, "_aigie_patched", False):
                original_stream = app.stream

                def traced_stream(input: Any, config: Dict[str, Any] | None = None, **kwargs):
                    from ..auto_instrument.trace import (
                        get_or_create_trace_sync,
                        set_active_callback_handler,
                        set_callback_context,
                        set_current_trace,
                    )
                    from ..callback import AigieCallbackHandler
                    from ..client import get_aigie

                    aigie = get_aigie()
                    if not aigie or not aigie._initialized:
                        yield from original_stream(input, config=config, **kwargs)
                        return

                    workflow_name = _extract_workflow_name(input, graph_schema)
                    trace = get_or_create_trace_sync(
                        name=workflow_name,
                        metadata={
                            "type": "langgraph",
                            "framework": "langgraph",
                            "inputs": input if isinstance(input, dict) else {},
                        },
                    )

                    handler = None
                    if trace:
                        set_current_trace(trace)
                        handler = AigieCallbackHandler(aigie=aigie, trace=trace)
                        handler._langgraph_mode = True
                        handler._workflow_name = workflow_name

                        if config is None:
                            config = {}
                        _safe_add_callback(config, handler)

                        if "metadata" not in config:
                            config["metadata"] = {}
                        config["metadata"]["run_name"] = workflow_name

                    set_callback_context(True)
                    set_active_callback_handler(handler)
                    try:
                        yield from original_stream(input, config=config, **kwargs)
                    finally:
                        set_active_callback_handler(None)
                        set_callback_context(False)

                traced_stream._aigie_patched = True
                app.stream = traced_stream

            return app

        traced_compile._aigie_patched = True
        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}")


_callback_manager_patched = False


def _patch_callback_manager() -> None:
    """Patch BaseCallbackManager.__init__ to auto-inject AigieCallbackHandler.

    Follows the OpenLLMetry pattern: when a LangGraph workflow is active, every
    new BaseCallbackManager automatically gets the workflow's handler injected.
    This ensures LLM calls inside graph nodes get proper parent-child span
    hierarchy even when the node doesn't forward config to the LLM.
    """
    global _callback_manager_patched
    if _callback_manager_patched:
        return

    try:
        from langchain_core.callbacks import BaseCallbackManager

        original_init = BaseCallbackManager.__init__

        @functools.wraps(original_init)
        def patched_init(self, *args, **kwargs):
            original_init(self, *args, **kwargs)

            # If a LangGraph workflow is active, inject the handler
            from ..auto_instrument.trace import get_active_callback_handler

            active_handler = get_active_callback_handler()
            if active_handler is None:
                return

            # Don't inject if this manager already has an AigieCallbackHandler
            for h in self.inheritable_handlers:
                if type(h).__name__ == "AigieCallbackHandler":
                    return

            self.add_handler(active_handler, True)

        patched_init._aigie_patched = True
        BaseCallbackManager.__init__ = patched_init
        _callback_manager_patched = True
        logger.debug("Patched BaseCallbackManager.__init__ for handler auto-injection")

    except ImportError:
        pass  # langchain-core not installed
    except Exception as e:
        logger.debug(f"Failed to patch BaseCallbackManager.__init__: {e}")
