from contextvars import copy_context
from importlib import import_module
from types import ModuleType
from typing import Any, Callable, Collection, Optional

from opentelemetry import trace as trace_api
from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore
from wrapt import wrap_function_wrapper

from openinference.instrumentation import (
    OITracer,
    TraceConfig,
)
from openinference.instrumentation.smolagents._wrappers import (
    _ModelStreamWrapper,
    _ModelWrapper,
    _RunWrapper,
    _StepWrapper,
    _ToolCallWrapper,
)
from openinference.instrumentation.smolagents.version import __version__

_instruments = ("smolagents >= 1.2.2.dev0",)


class SmolagentsInstrumentor(BaseInstrumentor):  # type: ignore
    __slots__ = (
        "_original_run_method",
        "_original_step_stream_methods",
        "_original_tool_call_method",
        "_original_model_generate_methods",
        "_original_model_generate_stream_methods",
        "_original_executor",
        "_tracer",
    )
    _original_executor: Optional[type]

    def instrumentation_dependencies(self) -> Collection[str]:
        return _instruments

    def _instrument(self, **kwargs: Any) -> None:
        import smolagents
        from smolagents import CodeAgent, MultiStepAgent, Tool, ToolCallingAgent, models

        if not (tracer_provider := kwargs.get("tracer_provider")):
            tracer_provider = trace_api.get_tracer_provider()
        if not (config := kwargs.get("config")):
            config = TraceConfig()
        else:
            assert isinstance(config, TraceConfig)
        self._tracer = OITracer(
            trace_api.get_tracer(__name__, __version__, tracer_provider),
            config=config,
        )

        run_wrapper = _RunWrapper(tracer=self._tracer)  # type: ignore[arg-type]
        self._original_run_method = getattr(MultiStepAgent, "run", None)
        wrap_function_wrapper(
            "smolagents",
            "MultiStepAgent.run",
            run_wrapper,
        )

        self._original_step_stream_methods: Optional[dict[type, Optional[Callable[..., Any]]]] = {}
        step_wrapper = _StepWrapper(tracer=self._tracer)  # type: ignore[arg-type]
        for step_cls in [CodeAgent, ToolCallingAgent]:
            self._original_step_stream_methods[step_cls] = getattr(step_cls, "_step_stream", None)
            wrap_function_wrapper(
                "smolagents",
                f"{step_cls.__name__}._step_stream",
                step_wrapper,
            )

        self._original_model_generate_methods: Optional[dict[type, Callable[..., Any]]] = {}
        self._original_model_generate_stream_methods: Optional[dict[type, Callable[..., Any]]] = {}

        exported_model_subclasses = list(
            dict.fromkeys(
                attr
                for _, attr in vars(smolagents).items()
                if isinstance(attr, type) and issubclass(attr, models.Model)
            )
        )

        for model_subclass in exported_model_subclasses:
            model_subclass_wrapper = _ModelWrapper(tracer=self._tracer)  # type: ignore[arg-type]

            self._original_model_generate_methods[model_subclass] = getattr(
                model_subclass, "generate"
            )
            wrap_function_wrapper(
                "smolagents",
                model_subclass.__name__ + ".generate",
                model_subclass_wrapper,
            )
            if original_generate_stream := model_subclass.__dict__.get("generate_stream"):
                self._original_model_generate_stream_methods[model_subclass] = (
                    original_generate_stream
                )
                wrap_function_wrapper(
                    "smolagents",
                    model_subclass.__name__ + ".generate_stream",
                    _ModelStreamWrapper(tracer=self._tracer),  # type: ignore[arg-type]
                )
        self._original_executor: Optional[type] = None

        module: ModuleType = import_module("smolagents.local_python_executor")

        _OriginalThreadPoolExecutor: type = getattr(module, "ThreadPoolExecutor")
        self._original_executor = _OriginalThreadPoolExecutor

        def _make_context_aware_executor(*args: Any, **_kwargs: Any) -> Any:
            executor = _OriginalThreadPoolExecutor(*args, **_kwargs)
            original_submit = executor.submit

            def context_preserving_submit(
                fn: Callable[..., Any], *fn_args: Any, **fn_kwargs: Any
            ) -> Any:
                ctx = copy_context()
                return original_submit(lambda: ctx.run(fn, *fn_args, **fn_kwargs))

            executor.submit = context_preserving_submit
            return executor

        setattr(module, "ThreadPoolExecutor", _make_context_aware_executor)

        tool_call_wrapper = _ToolCallWrapper(tracer=self._tracer)  # type: ignore[arg-type]
        self._original_tool_call_method = getattr(Tool, "__call__", None)
        wrap_function_wrapper(
            "smolagents",
            "Tool.__call__",
            tool_call_wrapper,
        )

    def _uninstrument(self, **kwargs: Any) -> None:
        from smolagents import MultiStepAgent, Tool

        if self._original_run_method is not None:
            MultiStepAgent.run = self._original_run_method
            self._original_run_method = None

        if self._original_step_stream_methods is not None:
            for step_cls, original_step_method in self._original_step_stream_methods.items():
                setattr(step_cls, "_step_stream", original_step_method)
            self._original_step_stream_methods = None

        if self._original_model_generate_methods is not None:
            for (
                model_subclass,
                original_model_generate_method,
            ) in self._original_model_generate_methods.items():
                setattr(model_subclass, "generate", original_model_generate_method)
            self._original_model_generate_methods = None

        if self._original_model_generate_stream_methods is not None:
            for (
                model_subclass,
                original_model_generate_stream_method,
            ) in self._original_model_generate_stream_methods.items():
                setattr(
                    model_subclass,
                    "generate_stream",
                    original_model_generate_stream_method,
                )
            self._original_model_generate_stream_methods = None

        if self._original_tool_call_method is not None:
            Tool.__call__ = self._original_tool_call_method
            self._original_tool_call_method = None

        if self._original_executor is not None:
            module: ModuleType = import_module("smolagents.local_python_executor")
            setattr(module, "ThreadPoolExecutor", self._original_executor)
            self._original_executor = None
