"""Tests for the patcher registry, Anthropic patcher, and OpenAI patcher."""

import asyncio
import sys
from types import SimpleNamespace, ModuleType
from typing import List, Sequence
from unittest.mock import MagicMock

import pytest
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider, ReadableSpan
from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter, SpanExportResult
from opentelemetry.trace import StatusCode

from pendo_sdk.patchers.registry import (
    _clear_registry,
    _registry,
    _active_unpatchers,
    register_patcher,
    patch_all,
    unpatch_all,
    get_registry,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

class _InMemoryExporter(SpanExporter):
    def __init__(self) -> None:
        self._spans: List[ReadableSpan] = []

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        self._spans.extend(spans)
        return SpanExportResult.SUCCESS

    def get_finished_spans(self) -> List[ReadableSpan]:
        return list(self._spans)

    def shutdown(self) -> None:
        pass


@pytest.fixture(autouse=True)
def _otel_setup():
    """Set up in-memory exporter and clean registry for each test."""
    _clear_registry()
    exporter = _InMemoryExporter()
    provider = TracerProvider()
    provider.add_span_processor(SimpleSpanProcessor(exporter))
    trace.set_tracer_provider(provider)
    yield exporter
    _clear_registry()
    provider.shutdown()
    trace._TRACER_PROVIDER = None
    trace._TRACER_PROVIDER_SET_ONCE._done = False


# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------

class TestRegistry:
    def test_register_and_get(self):
        def my_patcher():
            return None

        register_patcher("test_lib", my_patcher)
        reg = get_registry()
        assert "test_lib" in reg
        assert reg["test_lib"] is my_patcher

    def test_patch_all_calls_patchers(self):
        called = []

        def patcher_a():
            called.append("a")
            return lambda: None

        def patcher_b():
            called.append("b")
            return lambda: None

        register_patcher("a", patcher_a)
        register_patcher("b", patcher_b)

        patched = patch_all()
        assert set(patched) == {"a", "b"}
        assert set(called) == {"a", "b"}

    def test_patch_all_skips_if_returns_none(self):
        def optional_patcher():
            return None  # lib not installed

        register_patcher("optional", optional_patcher)
        patched = patch_all()
        assert patched == []
        assert "optional" not in _active_unpatchers

    def test_patch_all_skips_already_active(self):
        call_count = 0

        def counting_patcher():
            nonlocal call_count
            call_count += 1
            return lambda: None

        register_patcher("counter", counting_patcher)
        patch_all()
        patch_all()  # second call should skip
        assert call_count == 1

    def test_unpatch_all_calls_unpatchers(self):
        unpatched = []

        def patcher():
            return lambda: unpatched.append(True)

        register_patcher("lib", patcher)
        patch_all()
        unpatch_all()
        assert len(unpatched) == 1
        assert len(_active_unpatchers) == 0

    def test_patch_all_handles_patcher_exception(self):
        def broken_patcher():
            raise RuntimeError("oops")

        def good_patcher():
            return lambda: None

        register_patcher("broken", broken_patcher)
        register_patcher("good", good_patcher)

        patched = patch_all()
        assert "good" in patched
        assert "broken" not in patched


# ---------------------------------------------------------------------------
# Anthropic Patcher (with fake module)
# ---------------------------------------------------------------------------

class TestAnthropicPatcher:
    def _setup_fake_anthropic(self):
        """Create a fake anthropic module structure for patching."""
        # Build fake module hierarchy
        messages_mod = ModuleType("anthropic.resources.messages")

        class FakeMessages:
            def create(self, **kwargs):
                return SimpleNamespace(
                    model="claude-sonnet-4-20250514",
                    usage=SimpleNamespace(input_tokens=100, output_tokens=50),
                )

        class FakeAsyncMessages:
            async def create(self, **kwargs):
                return SimpleNamespace(
                    model="claude-sonnet-4-20250514",
                    usage=SimpleNamespace(input_tokens=200, output_tokens=80),
                )

        messages_mod.Messages = FakeMessages
        messages_mod.AsyncMessages = FakeAsyncMessages

        # Register in sys.modules so import works
        anthropic_mod = ModuleType("anthropic")
        resources_mod = ModuleType("anthropic.resources")
        sys.modules["anthropic"] = anthropic_mod
        sys.modules["anthropic.resources"] = resources_mod
        sys.modules["anthropic.resources.messages"] = messages_mod

        return FakeMessages, FakeAsyncMessages, messages_mod

    def _cleanup_fake_anthropic(self):
        for key in list(sys.modules):
            if key.startswith("anthropic"):
                del sys.modules[key]

    def test_patch_wraps_sync_create(self, _otel_setup):
        exporter = _otel_setup
        FakeMessages, _, _ = self._setup_fake_anthropic()
        try:
            from pendo_sdk.patchers.anthropic import patch_anthropic

            unpatch = patch_anthropic()
            assert unpatch is not None

            # Call the patched method
            client = FakeMessages()
            result = client.create(model="claude-sonnet-4-20250514", messages=[])

            assert result.model == "claude-sonnet-4-20250514"

            spans = exporter.get_finished_spans()
            assert len(spans) == 1
            span = spans[0]
            assert span.name == "anthropic.messages.create"
            assert span.attributes["pendo.span.type"] == "GENERATION"
            assert span.attributes["llm.model_name"] == "claude-sonnet-4-20250514"
            assert span.attributes["llm.token_count.prompt"] == 100
            assert span.attributes["llm.token_count.completion"] == 50

            unpatch()
        finally:
            self._cleanup_fake_anthropic()

    def test_patch_wraps_async_create(self, _otel_setup):
        exporter = _otel_setup
        _, FakeAsyncMessages, _ = self._setup_fake_anthropic()
        try:
            from pendo_sdk.patchers.anthropic import patch_anthropic

            unpatch = patch_anthropic()

            client = FakeAsyncMessages()
            result = asyncio.run(client.create(model="claude-sonnet-4-20250514", messages=[]))

            assert result.model == "claude-sonnet-4-20250514"

            spans = exporter.get_finished_spans()
            assert len(spans) == 1
            assert spans[0].attributes["llm.token_count.prompt"] == 200
            assert spans[0].attributes["llm.token_count.completion"] == 80

            unpatch()
        finally:
            self._cleanup_fake_anthropic()

    def test_unpatch_restores_original(self, _otel_setup):
        FakeMessages, _, messages_mod = self._setup_fake_anthropic()
        try:
            original_create = FakeMessages.create
            from pendo_sdk.patchers.anthropic import patch_anthropic

            unpatch = patch_anthropic()
            assert FakeMessages.create is not original_create

            unpatch()
            assert FakeMessages.create is original_create
        finally:
            self._cleanup_fake_anthropic()

    def test_returns_none_when_not_installed(self):
        import importlib
        import builtins

        original_import = builtins.__import__

        def mock_import(name, *args, **kwargs):
            if name.startswith("anthropic"):
                raise ImportError("mocked: anthropic not installed")
            return original_import(name, *args, **kwargs)

        builtins.__import__ = mock_import
        try:
            # Re-import to get a fresh reference
            from pendo_sdk.patchers.anthropic import patch_anthropic
            result = patch_anthropic()
            assert result is None
        finally:
            builtins.__import__ = original_import

    def test_exception_sets_error_status(self, _otel_setup):
        exporter = _otel_setup
        messages_mod = ModuleType("anthropic.resources.messages")

        class ErrorMessages:
            def create(self, **kwargs):
                raise ConnectionError("API down")

        messages_mod.Messages = ErrorMessages
        messages_mod.AsyncMessages = None

        sys.modules["anthropic"] = ModuleType("anthropic")
        sys.modules["anthropic.resources"] = ModuleType("anthropic.resources")
        sys.modules["anthropic.resources.messages"] = messages_mod

        try:
            from pendo_sdk.patchers.anthropic import patch_anthropic

            unpatch = patch_anthropic()
            client = ErrorMessages()

            with pytest.raises(ConnectionError, match="API down"):
                client.create(model="test")

            spans = exporter.get_finished_spans()
            assert len(spans) == 1
            assert spans[0].status.status_code == StatusCode.ERROR

            if unpatch:
                unpatch()
        finally:
            self._cleanup_fake_anthropic()


# ---------------------------------------------------------------------------
# OpenAI Patcher (with fake module)
# ---------------------------------------------------------------------------

class TestOpenAIPatcher:
    def _setup_fake_openai(self):
        """Create a fake openai module structure for patching."""
        completions_mod = ModuleType("openai.resources.chat.completions")

        class FakeCompletions:
            def create(self, **kwargs):
                return SimpleNamespace(
                    model="gpt-4",
                    usage=SimpleNamespace(
                        prompt_tokens=150,
                        completion_tokens=60,
                        total_tokens=210,
                    ),
                )

        class FakeAsyncCompletions:
            async def create(self, **kwargs):
                return SimpleNamespace(
                    model="gpt-4",
                    usage=SimpleNamespace(
                        prompt_tokens=300,
                        completion_tokens=120,
                        total_tokens=420,
                    ),
                )

        completions_mod.Completions = FakeCompletions
        completions_mod.AsyncCompletions = FakeAsyncCompletions

        openai_mod = ModuleType("openai")
        resources_mod = ModuleType("openai.resources")
        chat_mod = ModuleType("openai.resources.chat")
        sys.modules["openai"] = openai_mod
        sys.modules["openai.resources"] = resources_mod
        sys.modules["openai.resources.chat"] = chat_mod
        sys.modules["openai.resources.chat.completions"] = completions_mod

        return FakeCompletions, FakeAsyncCompletions, completions_mod

    def _cleanup_fake_openai(self):
        for key in list(sys.modules):
            if key.startswith("openai"):
                del sys.modules[key]

    def test_patch_wraps_sync_create(self, _otel_setup):
        exporter = _otel_setup
        FakeCompletions, _, _ = self._setup_fake_openai()
        try:
            from pendo_sdk.patchers.openai import patch_openai

            unpatch = patch_openai()
            assert unpatch is not None

            client = FakeCompletions()
            result = client.create(model="gpt-4", messages=[])

            assert result.model == "gpt-4"

            spans = exporter.get_finished_spans()
            assert len(spans) == 1
            span = spans[0]
            assert span.name == "openai.chat.completions.create"
            assert span.attributes["pendo.span.type"] == "GENERATION"
            assert span.attributes["llm.model_name"] == "gpt-4"
            assert span.attributes["llm.token_count.prompt"] == 150
            assert span.attributes["llm.token_count.completion"] == 60
            assert span.attributes["llm.token_count.total"] == 210

            unpatch()
        finally:
            self._cleanup_fake_openai()

    def test_patch_wraps_async_create(self, _otel_setup):
        exporter = _otel_setup
        _, FakeAsyncCompletions, _ = self._setup_fake_openai()
        try:
            from pendo_sdk.patchers.openai import patch_openai

            unpatch = patch_openai()

            client = FakeAsyncCompletions()
            result = asyncio.run(client.create(model="gpt-4", messages=[]))

            assert result.model == "gpt-4"

            spans = exporter.get_finished_spans()
            assert len(spans) == 1
            assert spans[0].attributes["llm.token_count.prompt"] == 300
            assert spans[0].attributes["llm.token_count.total"] == 420

            unpatch()
        finally:
            self._cleanup_fake_openai()

    def test_unpatch_restores_original(self, _otel_setup):
        FakeCompletions, _, _ = self._setup_fake_openai()
        try:
            original_create = FakeCompletions.create
            from pendo_sdk.patchers.openai import patch_openai

            unpatch = patch_openai()
            assert FakeCompletions.create is not original_create

            unpatch()
            assert FakeCompletions.create is original_create
        finally:
            self._cleanup_fake_openai()

    def test_returns_none_when_not_installed(self):
        for key in list(sys.modules):
            if key.startswith("openai"):
                del sys.modules[key]

        from pendo_sdk.patchers.openai import patch_openai
        result = patch_openai()
        assert result is None

    def test_exception_sets_error_status(self, _otel_setup):
        exporter = _otel_setup
        completions_mod = ModuleType("openai.resources.chat.completions")

        class ErrorCompletions:
            def create(self, **kwargs):
                raise TimeoutError("request timed out")

        completions_mod.Completions = ErrorCompletions
        completions_mod.AsyncCompletions = None

        sys.modules["openai"] = ModuleType("openai")
        sys.modules["openai.resources"] = ModuleType("openai.resources")
        sys.modules["openai.resources.chat"] = ModuleType("openai.resources.chat")
        sys.modules["openai.resources.chat.completions"] = completions_mod

        try:
            from pendo_sdk.patchers.openai import patch_openai

            unpatch = patch_openai()
            client = ErrorCompletions()

            with pytest.raises(TimeoutError, match="request timed out"):
                client.create(model="gpt-4")

            spans = exporter.get_finished_spans()
            assert len(spans) == 1
            assert spans[0].status.status_code == StatusCode.ERROR

            if unpatch:
                unpatch()
        finally:
            self._cleanup_fake_openai()


# ---------------------------------------------------------------------------
# Integration: registry + patchers together
# ---------------------------------------------------------------------------

class TestRegistryIntegration:
    def test_custom_patcher_registration(self, _otel_setup):
        """Users can register their own patchers for custom LLM providers."""
        exporter = _otel_setup
        patched_libs = []

        def custom_bedrock_patcher():
            patched_libs.append("bedrock")
            return lambda: patched_libs.remove("bedrock")

        register_patcher("bedrock", custom_bedrock_patcher)
        patched = patch_all()

        assert "bedrock" in patched
        assert "bedrock" in patched_libs

        unpatch_all()
        assert "bedrock" not in patched_libs
