"""
Tests for OpenAI / Anthropic / LangChain interceptors.

Uses pure mocking — no API keys or network calls needed.
Proves:
  1. Interceptors start/stop cleanly.
  2. Events are emitted with correct fields on a successful call.
  3. Events are emitted on errors (with success=False).
  4. Double-counting guard works: nested calls emit only one event.
  5. Zero-SDK tracker init still succeeds.
  6. Dynamic provider normalization works without a hardcoded map.

Run with: pytest tests/test_interceptors.py -v
"""

import sys
import time
import types
import threading
import pytest
from unittest.mock import MagicMock, patch, PropertyMock


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

def _make_usage(prompt_tokens=10, completion_tokens=20):
    """Mock matching OpenAI's response.usage shape."""
    usage = MagicMock()
    usage.prompt_tokens = prompt_tokens
    usage.completion_tokens = completion_tokens
    return usage


def _make_anthropic_usage(input_tokens=10, output_tokens=20):
    """Mock matching Anthropic's response.usage shape."""
    usage = MagicMock()
    usage.input_tokens = input_tokens
    usage.output_tokens = output_tokens
    return usage


def _set_config():
    from agentcost.config import AgentCostConfig, set_config
    config = AgentCostConfig(api_key="test-key", project_id="test-project", enabled=True)
    set_config(config)
    return config


def _clear_config():
    from agentcost.config import set_config
    set_config(None)


# ===========================================================================
# OpenAI Interceptor Tests
# ===========================================================================

class TestOpenAIInterceptor:

    def test_start_stop_lifecycle(self):
        """Interceptor starts, patches Completions.create, then restores it."""
        from agentcost.openai_interceptor import OpenAIInterceptor

        interceptor = OpenAIInterceptor(event_callback=lambda e: None)

        ok = interceptor.start()
        assert ok is True
        assert interceptor.is_active is True

        interceptor.stop()
        assert interceptor.is_active is False

    def test_idempotent_start(self):
        """Calling start() twice is safe."""
        from agentcost.openai_interceptor import OpenAIInterceptor

        interceptor = OpenAIInterceptor(event_callback=lambda e: None)
        first = interceptor.start()
        second = interceptor.start()
        assert first is True
        assert second is True
        assert interceptor.is_active is True
        interceptor.stop()

    def test_event_emitted_on_call(self):
        """Patched create() emits exactly one event with correct fields."""
        from agentcost.openai_interceptor import OpenAIInterceptor, _tracking_depth
        from openai.resources.chat.completions import Completions

        _tracking_depth.value = 0
        _set_config()

        # Save real method, put a stub as "original" before starting interceptor
        real_create = Completions.create

        fake_response = MagicMock()
        fake_response.usage = _make_usage(prompt_tokens=50, completion_tokens=100)

        def stub_create(self_arg, *args, **kwargs):
            return fake_response

        Completions.create = stub_create  # interceptor will save this as _original

        events = []
        interceptor = OpenAIInterceptor(event_callback=lambda e: events.append(e))
        interceptor.start()

        try:
            dummy_client = MagicMock()
            result = Completions.create(
                dummy_client,
                model="gpt-4o",
                messages=[{"role": "user", "content": "Hello"}],
            )

            assert result is fake_response
            assert len(events) == 1

            evt = events[0]
            assert evt["model"] == "gpt-4o"
            assert evt["input_tokens"] == 50
            assert evt["output_tokens"] == 100
            assert evt["success"] is True
            assert evt["error"] is None
            assert "cost" in evt
            assert "latency_ms" in evt
            assert "input_hash" in evt
            assert "timestamp" in evt
        finally:
            interceptor.stop()
            Completions.create = real_create
            _clear_config()

    def test_event_on_error(self):
        """An exception in create() still emits an event with success=False."""
        from agentcost.openai_interceptor import OpenAIInterceptor, _tracking_depth
        from openai.resources.chat.completions import Completions

        _tracking_depth.value = 0
        _set_config()

        real_create = Completions.create

        def boom_create(self_arg, *args, **kwargs):
            raise RuntimeError("API rate limit")

        Completions.create = boom_create
        events = []
        interceptor = OpenAIInterceptor(event_callback=lambda e: events.append(e))
        interceptor.start()

        try:
            dummy = MagicMock()
            with pytest.raises(RuntimeError, match="API rate limit"):
                Completions.create(
                    dummy,
                    model="gpt-4o",
                    messages=[{"role": "user", "content": "test"}],
                )

            assert len(events) == 1
            assert events[0]["success"] is False
            assert "API rate limit" in events[0]["error"]
        finally:
            interceptor.stop()
            Completions.create = real_create
            _clear_config()

    def test_depth_guard_prevents_double_count(self):
        """When _tracking_depth > 0 the interceptor skips entirely."""
        from agentcost.openai_interceptor import OpenAIInterceptor, _tracking_depth
        from openai.resources.chat.completions import Completions

        _set_config()

        real_create = Completions.create

        fake_response = MagicMock()
        fake_response.usage = _make_usage(10, 20)

        call_count = {"n": 0}

        def stub_create(self_arg, *args, **kwargs):
            call_count["n"] += 1
            return fake_response

        Completions.create = stub_create
        events = []
        interceptor = OpenAIInterceptor(event_callback=lambda e: events.append(e))
        interceptor.start()

        try:
            # Simulate being inside a LangChain interceptor
            _tracking_depth.value = 1

            dummy = MagicMock()
            result = Completions.create(
                dummy,
                model="gpt-4o",
                messages=[{"role": "user", "content": "nested"}],
            )

            assert result is fake_response
            assert call_count["n"] == 1   # Original was still called
            assert len(events) == 0       # But NO event was emitted
        finally:
            _tracking_depth.value = 0
            interceptor.stop()
            Completions.create = real_create
            _clear_config()


# ===========================================================================
# Anthropic Interceptor Tests
# ===========================================================================

class TestAnthropicInterceptor:

    def test_start_stop_lifecycle(self):
        """Interceptor starts, patches Messages.create, then restores it."""
        from agentcost.anthropic_interceptor import AnthropicInterceptor

        interceptor = AnthropicInterceptor(event_callback=lambda e: None)

        ok = interceptor.start()
        assert ok is True
        assert interceptor.is_active is True

        interceptor.stop()
        assert interceptor.is_active is False

    def test_event_emitted_on_call(self):
        """Patched Messages.create() emits exactly one event."""
        from agentcost.anthropic_interceptor import AnthropicInterceptor, _tracking_depth
        from anthropic.resources.messages import Messages

        _tracking_depth.value = 0
        _set_config()

        real_create = Messages.create
        fake_response = MagicMock()
        fake_response.usage = _make_anthropic_usage(input_tokens=30, output_tokens=60)

        def stub_create(self_arg, *args, **kwargs):
            return fake_response

        Messages.create = stub_create
        events = []
        interceptor = AnthropicInterceptor(event_callback=lambda e: events.append(e))
        interceptor.start()

        try:
            dummy = MagicMock()
            result = Messages.create(
                dummy,
                model="claude-3-5-sonnet-20241022",
                messages=[{"role": "user", "content": "Hello"}],
            )

            assert result is fake_response
            assert len(events) == 1

            evt = events[0]
            assert evt["model"] == "claude-3-5-sonnet-20241022"
            assert evt["input_tokens"] == 30
            assert evt["output_tokens"] == 60
            assert evt["success"] is True
        finally:
            interceptor.stop()
            Messages.create = real_create
            _clear_config()

    def test_depth_guard_prevents_double_count(self):
        """Anthropic interceptor skips when depth > 0."""
        from agentcost.anthropic_interceptor import AnthropicInterceptor, _tracking_depth
        from anthropic.resources.messages import Messages

        _set_config()

        real_create = Messages.create
        fake_response = MagicMock()
        fake_response.usage = _make_anthropic_usage(10, 20)

        call_count = {"n": 0}

        def stub_create(self_arg, *args, **kwargs):
            call_count["n"] += 1
            return fake_response

        Messages.create = stub_create
        events = []
        interceptor = AnthropicInterceptor(event_callback=lambda e: events.append(e))
        interceptor.start()

        try:
            _tracking_depth.value = 1

            dummy = MagicMock()
            result = Messages.create(
                dummy,
                model="claude-3-5-sonnet-20241022",
                messages=[{"role": "user", "content": "nested"}],
            )

            assert result is fake_response
            assert call_count["n"] == 1
            assert len(events) == 0
        finally:
            _tracking_depth.value = 0
            interceptor.stop()
            Messages.create = real_create
            _clear_config()


# ===========================================================================
# LangChain Interceptor + Depth Guard Integration
# ===========================================================================

class TestLangChainDepthGuard:
    """Verify the LangChain interceptor sets the depth guard so nested SDK calls skip."""

    def test_langchain_sets_depth_during_invoke(self):
        """While LangChain invoke runs, _tracking_depth > 0 — blocking nested events."""
        from agentcost.interceptor import LangChainInterceptor
        try:
            from agentcost.openai_interceptor import _tracking_depth
        except ImportError:
            _tracking_depth = threading.local()

        from langchain_core.language_models import BaseChatModel

        _tracking_depth.value = 0
        _set_config()

        real_invoke = BaseChatModel.invoke

        depth_during_call = []

        # This stub simulates what the LLM does under the hood.
        # Inside here, a real app would call openai/anthropic SDK.
        # We capture the depth to verify it was incremented.
        def fake_invoke(llm_self, input_data, *args, **kwargs):
            depth_during_call.append(getattr(_tracking_depth, 'value', 0))
            result = MagicMock()
            result.content = "response text"
            return result

        BaseChatModel.invoke = fake_invoke

        events = []
        interceptor = LangChainInterceptor(event_callback=lambda e: events.append(e))
        interceptor.start()

        try:
            llm = MagicMock(spec=BaseChatModel)
            llm.model_name = "gpt-4o"
            type(llm)._llm_type = PropertyMock(return_value="chat_openai")

            BaseChatModel.invoke(llm, "test input")

            # During the call, depth was > 0
            assert len(depth_during_call) == 1
            assert depth_during_call[0] > 0

            # After the call, depth is back to 0
            assert getattr(_tracking_depth, 'value', 0) == 0

            # LangChain interceptor emitted exactly one event
            assert len(events) == 1
            assert events[0]["model"] == "gpt-4o"
        finally:
            interceptor.stop()
            BaseChatModel.invoke = real_invoke
            _clear_config()

    def test_end_to_end_no_double_count(self):
        """LangChain invoke calling OpenAI create emits exactly 1 event total."""
        from agentcost.interceptor import LangChainInterceptor
        from agentcost.openai_interceptor import OpenAIInterceptor, _tracking_depth
        from langchain_core.language_models import BaseChatModel
        from openai.resources.chat.completions import Completions

        _tracking_depth.value = 0
        _set_config()

        real_invoke = BaseChatModel.invoke
        real_create = Completions.create

        all_events = []

        def collector(e):
            all_events.append(e)

        # Fake response from "OpenAI"
        fake_oai_response = MagicMock()
        fake_oai_response.usage = _make_usage(50, 100)

        def stub_oai_create(self_arg, *args, **kwargs):
            return fake_oai_response

        Completions.create = stub_oai_create

        # LangChain's invoke calls OpenAI under the hood
        def fake_invoke(llm_self, input_data, *args, **kwargs):
            dummy = MagicMock()
            Completions.create(
                dummy,
                model="gpt-4o",
                messages=[{"role": "user", "content": str(input_data)}],
            )
            result = MagicMock()
            result.content = "answer"
            return result

        BaseChatModel.invoke = fake_invoke

        # Start BOTH interceptors (simulates real tracker behavior)
        lc_interceptor = LangChainInterceptor(event_callback=collector)
        oai_interceptor = OpenAIInterceptor(event_callback=collector)
        lc_interceptor.start()
        oai_interceptor.start()

        try:
            llm = MagicMock(spec=BaseChatModel)
            llm.model_name = "gpt-4o"
            type(llm)._llm_type = PropertyMock(return_value="chat_openai")

            BaseChatModel.invoke(llm, "Tell me a joke")

            # THE KEY ASSERTION: only 1 event, not 2
            assert len(all_events) == 1, (
                f"Expected 1 event but got {len(all_events)}. "
                f"Double-counting guard failed!"
            )
        finally:
            oai_interceptor.stop()
            lc_interceptor.stop()
            BaseChatModel.invoke = real_invoke
            Completions.create = real_create
            _tracking_depth.value = 0
            _clear_config()


# ===========================================================================
# Zero-SDK Tracker Test
# ===========================================================================

class TestZeroSDKInit:
    """Verify tracker init succeeds even when no LLM SDK is installed."""

    def test_tracker_init_with_no_sdks(self):
        """Tracker becomes initialized even if all interceptors fail to start."""
        from agentcost.tracker import AgentCostTracker

        tracker = AgentCostTracker()

        # Create mock interceptor classes whose start() always returns False
        class FakeInterceptor:
            def __init__(self, **kwargs):
                pass
            def start(self):
                return False

        # Patch the import locations so the tracker gets our failing classes
        with patch("agentcost.interceptor.LangChainInterceptor", FakeInterceptor), \
             patch("agentcost.openai_interceptor.OpenAIInterceptor", FakeInterceptor), \
             patch("agentcost.anthropic_interceptor.AnthropicInterceptor", FakeInterceptor):
            tracker.init(
                api_key="test-key",
                project_id="test-project",
                enabled=True,
                local_mode=True,
            )

        # Even with no interceptors active, tracker should be initialized
        assert tracker._is_initialized is True
        assert tracker.is_active is True
        assert len(tracker._interceptors) == 0
        tracker.shutdown()
        _clear_config()


# ===========================================================================
# Provider Normalization Tests (Backend)
# ===========================================================================

class TestDynamicProviderNormalization:
    """Verify _normalize_provider works dynamically, not from a hardcoded map."""

    @pytest.fixture(autouse=True)
    def add_backend_to_path(self):
        sys.path.insert(0, "d:\\AgentCost\\agentcost-backend")
        yield
        if "d:\\AgentCost\\agentcost-backend" in sys.path:
            sys.path.remove("d:\\AgentCost\\agentcost-backend")

    def _get_service(self):
        from app.services.pricing_service import PricingService
        return PricingService(db=MagicMock())

    def test_known_platforms_remap(self):
        svc = self._get_service()
        assert svc._normalize_provider("vertex_ai") == "google"
        assert svc._normalize_provider("vertex_ai-anthropic_models") == "google"
        assert svc._normalize_provider("bedrock") == "aws"
        assert svc._normalize_provider("bedrock_converse") == "aws"
        assert svc._normalize_provider("azure") == "azure"
        assert svc._normalize_provider("azure_ai") == "azure"
        assert svc._normalize_provider("text-completion-openai") == "openai"
        assert svc._normalize_provider("palm") == "google"
        assert svc._normalize_provider("gemini") == "google"
        assert svc._normalize_provider("watsonx") == "ibm"
        assert svc._normalize_provider("oci") == "oracle"

    def test_dynamic_suffix_stripping(self):
        svc = self._get_service()
        assert svc._normalize_provider("together_ai") == "together"
        assert svc._normalize_provider("fireworks_ai") == "fireworks"
        assert svc._normalize_provider("jina_ai") == "jina"

    def test_simple_providers_pass_through(self):
        svc = self._get_service()
        assert svc._normalize_provider("openai") == "openai"
        assert svc._normalize_provider("anthropic") == "anthropic"
        assert svc._normalize_provider("groq") == "groq"
        assert svc._normalize_provider("mistral") == "mistral"
        assert svc._normalize_provider("deepseek") == "deepseek"
        assert svc._normalize_provider("cohere") == "cohere"

    def test_unknown_provider_passes_through(self):
        """A new provider LiteLLM adds tomorrow should NOT be dropped."""
        svc = self._get_service()
        assert svc._normalize_provider("somenewprovider") == "somenewprovider"
        assert svc._normalize_provider("FutureAI") == "futureai"

    def test_empty_and_whitespace(self):
        svc = self._get_service()
        assert svc._normalize_provider("") == "unknown"
        assert svc._normalize_provider("  ") == "unknown"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
