from autogen import ConversableAgent
from opentelemetry.util._importlib_metadata import entry_points

from openinference.instrumentation.ag2 import AG2Instrumentor
from openinference.instrumentation.autogen import AutogenInstrumentor, SpanAttributes


class TestInstrumentor:
    def test_entrypoint_for_opentelemetry_instrument(self) -> None:
        (instrumentor_entrypoint,) = entry_points(
            group="opentelemetry_instrumentor", name="autogen"
        )
        instrumentor = instrumentor_entrypoint.load()()
        assert isinstance(instrumentor, AutogenInstrumentor)

    def test_legacy_span_attributes_remain_importable(self) -> None:
        assert SpanAttributes.OPENINFERENCE_SPAN_KIND == "openinference.span.kind"
        assert SpanAttributes.TOOL_NAME == "tool.name"

    def test_uninstrument_ignores_instrumentation_it_does_not_own(self) -> None:
        original = ConversableAgent.generate_reply
        ag2_instrumentor = AG2Instrumentor()
        ag2_instrumentor.instrument()
        try:
            AutogenInstrumentor().uninstrument()
            assert ConversableAgent.generate_reply is not original
        finally:
            ag2_instrumentor.uninstrument()

        assert ConversableAgent.generate_reply is original

    def test_uninstrument_restores_methods_when_ag2_instrumented_first(self) -> None:
        original = ConversableAgent.generate_reply
        ag2_instrumentor = AG2Instrumentor()
        ag2_instrumentor.instrument()  # e.g. auto-instrumentation loading the ag2 entry point
        try:
            facade = AutogenInstrumentor()
            facade.instrument()
            facade.uninstrument()
            assert ConversableAgent.generate_reply is original
        finally:
            ag2_instrumentor.uninstrument()

    def test_legacy_distribution_delegates_to_ag2_instrumentor(self) -> None:
        original = ConversableAgent.generate_reply
        instrumentor = AutogenInstrumentor()

        instrumentor.instrument()
        try:
            agent = ConversableAgent(
                "assistant",
                llm_config=False,
                code_execution_config=False,
                human_input_mode="NEVER",
                default_auto_reply="done",
            )
            assert agent.generate_reply(messages=[{"role": "user", "content": "hello"}]) == "done"
            assert instrumentor.is_instrumented_by_opentelemetry
            assert AG2Instrumentor().is_instrumented_by_opentelemetry
            assert ConversableAgent.generate_reply is not original
        finally:
            instrumentor.uninstrument()

        assert ConversableAgent.generate_reply is original
