import logging
from typing import Any, Collection, List, Tuple

from opentelemetry import trace as trace_api
from opentelemetry.instrumentation.instrumentor import (  # type: ignore[attr-defined]
    BaseInstrumentor,
)
from wrapt import resolve_path, wrap_function_wrapper

from openinference.instrumentation import OITracer, TraceConfig
from openinference.instrumentation.autogen_agentchat.version import __version__

_instruments = ("autogen-agentchat >= 0.5.1",)

logger = logging.getLogger(__name__)


class AutogenAgentChatInstrumentor(BaseInstrumentor):  # type: ignore
    """An instrumentor for autogen-agentchat"""

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

    def _instrument(self, **kwargs: Any) -> None:
        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,
        )

        from autogen_agentchat.agents import AssistantAgent, BaseChatAgent
        from autogen_agentchat.teams import BaseGroupChat
        from openinference.instrumentation.autogen_agentchat._wrappers import (
            _AssistantAgentExecuteToolCallWrapper,
            _AssistantAgentOnMessagesStreamWrapper,
            _BaseChatAgentOnMessagesStreamWrapper,
            _BaseGroupChatRunStreamWrapper,
            _BaseOpenAIChatCompletionClientCreateStreamWrapper,
            _BaseOpenAIChatCompletionClientCreateWrapper,
        )

        self._originals: List[Tuple[Any, Any, Any]] = []

        method_wrappers: dict[Any, Any] = {
            AssistantAgent.on_messages_stream: _AssistantAgentOnMessagesStreamWrapper(self._tracer),  # type: ignore[arg-type]
            BaseChatAgent.on_messages_stream: _BaseChatAgentOnMessagesStreamWrapper(self._tracer),  # type: ignore[arg-type]
            BaseGroupChat.run_stream: _BaseGroupChatRunStreamWrapper(self._tracer),  # type: ignore[arg-type]
            AssistantAgent._execute_tool_call: _AssistantAgentExecuteToolCallWrapper(self._tracer),  # type: ignore[arg-type]
        }

        # `autogen_ext.models.openai` imports `openai`; keep non-OpenAI clients usable.
        try:
            from autogen_ext.models.openai import BaseOpenAIChatCompletionClient
        except ImportError:
            logger.debug(
                "autogen_ext.models.openai is unavailable (openai not installed); "
                "skipping OpenAI chat-completion client instrumentation."
            )
        else:
            method_wrappers[BaseOpenAIChatCompletionClient.create] = (
                _BaseOpenAIChatCompletionClientCreateWrapper(self._tracer)  # type: ignore[arg-type]
            )
            method_wrappers[BaseOpenAIChatCompletionClient.create_stream] = (
                _BaseOpenAIChatCompletionClientCreateStreamWrapper(self._tracer)  # type: ignore[arg-type]
            )

        for method, wrapper in method_wrappers.items():
            module, name = method.__module__, method.__qualname__
            self._originals.append(resolve_path(module, name))
            wrap_function_wrapper(module, name, wrapper)

    def _uninstrument(self, **kwargs: Any) -> None:
        for parent, attribute, original in getattr(self, "_originals", ()):
            setattr(parent, attribute, original)
