from __future__ import annotations

from ._base import BaseInstrumentor
from ._patch_utils import wrap_sync, wrap_async
from agenx import SpanType

import wrapt


class LangChainInstrumentor(BaseInstrumentor):
    required_packages = ["langchain_core>=1.0.0"]

    def _instrument(self):
        """Applying patches for SpanType:

        - STEP
        - LLM 
        - AGENT
        - RETRIEVER
        - TOOL
        """

        """Chain / LCEL pipelines"""

        wrap_sync(
            "langchain_core.runnables.base", "RunnableSequence.invoke",
            span_name="langchain.chain", span_type=SpanType.STEP,
            extract_input=_lc_extract_input, extract_output=_lc_extract_output,
        )
        wrap_async(
            "langchain_core.runnables.base", "RunnableSequence.ainvoke",
            span_name="langchain.chain", span_type=SpanType.STEP,
            extract_input=_lc_extract_input, extract_output=_lc_extract_output,
        )

        """Direct LLM calls"""

        wrap_sync(
            "langchain_core.language_models.chat_models", "BaseChatModel.invoke",
            span_name="langchain.llm", span_type=SpanType.LLM,
            extract_input=_lc_extract_llm_input, extract_output=_lc_extract_output,
            enrich_span=_enrich_llm_span,
        )
        wrap_async(
            "langchain_core.language_models.chat_models", "BaseChatModel.ainvoke",
            span_name="langchain.llm", span_type=SpanType.LLM,
            extract_input=_lc_extract_llm_input, extract_output=_lc_extract_output,
            enrich_span=_enrich_llm_span,
        )

        """Agent executor"""

        wrap_sync(
            "langgraph.graph.state",  "CompiledStateGraph.invoke",
            span_name="langchain.agent", span_type=SpanType.AGENT,
            extract_input=_lc_extract_input, extract_output=_lc_extract_output,
        )
        wrap_async(
            "langgraph.pregel.main", "Pregel.ainvoke",
            span_name="langchain.agent", span_type=SpanType.AGENT,
            extract_input=_lc_extract_input, extract_output=_lc_extract_output,
        )

        """Tool calls"""

        wrap_sync(
            "langchain_core.tools.base", "BaseTool.invoke",
            span_name="langchain.tool", span_type=SpanType.TOOL,
            extract_input=_lc_extract_tool_input, extract_output=_lc_extract_tool_output,
            enrich_span=_enrich_tool_span,
        )
        wrap_async(
            "langchain_core.tools.base", "BaseTool.ainvoke",
            span_name="langchain.tool", span_type=SpanType.TOOL,
            extract_input=_lc_extract_tool_input, extract_output=_lc_extract_tool_output,
            enrich_span=_enrich_tool_span,
        )

        wrap_sync(
            "langchain_core.tools.base", "BaseTool.run",
            span_name="langchain.tool", span_type=SpanType.TOOL,
            extract_input=_lc_extract_tool_input, extract_output=_lc_extract_tool_output,
            enrich_span=_enrich_tool_span,
        )
        wrap_async(
            "langchain_core.tools.base", "BaseTool.arun",
            span_name="langchain.tool", span_type=SpanType.TOOL,
            extract_input=_lc_extract_tool_input, extract_output=_lc_extract_tool_output,
            enrich_span=_enrich_tool_span,
        )

    def _uninstrument(self):
        from ._patch_utils import _patched

        targets = [
            ("langchain_core.runnables.base", "RunnableSequence.invoke"),
            ("langchain_core.runnables.base", "RunnableSequence.ainvoke"),
            ("langchain_core.language_models.chat_models", "BaseChatModel.invoke"),
            ("langchain_core.language_models.chat_models", "BaseChatModel.ainvoke"),
            ("langgraph.graph.state", "CompiledStateGraph.invoke"),
            ("langgraph.pregel.main", "Pregel.ainvoke"),
            ("langchain_core.tools.base", "BaseTool.invoke"),
            ("langchain_core.tools.base", "BaseTool.ainvoke"),
            ("langchain_core.tools.base", "BaseTool.run"),
            ("langchain_core.tools.base", "BaseTool.arun"),
        ]
        for mod_path, attr in targets:
            _patched.discard(f"{mod_path}.{attr}")
            try:
                wrapt.unpatch(mod_path, attr)
            except Exception:
                pass

def _lc_extract_input(args, kwargs, instance):
    return args[0] if args else kwargs.get("input")


def _lc_extract_output(result, instance):
    if hasattr(result, "content"):        
        return result.content
    if hasattr(result, "__dict__"):
        return result.__dict__
    return result


def _lc_extract_llm_input(args, kwargs, instance):
    messages = args[0] if args else kwargs.get("input", [])
    if isinstance(messages, list):
        return [m.content if hasattr(m, "content") else str(m) for m in messages]
    return str(messages)


def _enrich_llm_span(span, instance, args, kwargs):
    model = getattr(instance, "model_name", None) or getattr(instance, "model", None)
    if model:
        span.llm.model = model
    provider = type(instance).__module__.split(".")[0]  
    span.llm.provider = provider


def _lc_extract_tool_input(args, kwargs, instance):
    return args[0] if args else kwargs.get("input")


def _lc_extract_tool_output(result, instance):
    return str(result) if result is not None else None


def _enrich_tool_span(span, instance, args, kwargs):
    span.tool.tool_name = getattr(instance, "name", type(instance).__name__)



