"""Tests for amplitude_ai.integrations.opentelemetry."""
from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock, patch

import pytest

_MOD = "amplitude_ai.integrations.opentelemetry"


class TestAmplitudeAgentExporter:
    """Tests for inbound OTEL span -> Amplitude event mapping."""

    @pytest.fixture(autouse=True)
    def _patch_otel_symbols(self) -> Any:
        """Ensure OTEL module-level symbols are available regardless of install."""
        mock_result = MagicMock()
        mock_result.SUCCESS = "SUCCESS"
        with (
            patch(f"{_MOD}.OTEL_AVAILABLE", True),
            patch(f"{_MOD}.SpanExportResult", mock_result),
        ):
            yield

    def _make_span(
        self,
        attrs: dict[str, Any],
        name: str = "chat",
        start_ns: int = 1_000_000_000,
        end_ns: int = 2_000_000_000,
    ) -> MagicMock:
        """Create a mock OTEL ReadableSpan with the given attributes."""
        span = MagicMock()
        span.name = name
        span.attributes = attrs
        span.start_time = start_ns
        span.end_time = end_ns
        span.context = MagicMock()
        span.context.trace_id = 0x1234567890ABCDEF
        span.context.span_id = 0xFEDCBA0987654321
        span.instrumentation_scope = None
        return span

    def _make_exporter(self, user_id: str = "test-user") -> Any:
        """Create an AmplitudeAgentExporter with a mocked Amplitude client."""
        from amplitude_ai.integrations.opentelemetry import AmplitudeAgentExporter

        return AmplitudeAgentExporter(amplitude=MagicMock(), user_id=user_id)

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.calculate_cost", return_value=0.001)
    @patch(f"{_MOD}.track_ai_message")
    def test_model_priority_prefers_response_model(
        self, mock_track_ai: MagicMock, mock_cost: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """gen_ai.response.model should take precedence over gen_ai.request.model."""
        exporter = self._make_exporter()
        span = self._make_span({
            "gen_ai.operation.name": "chat",
            "gen_ai.request.model": "gpt-4o",
            "gen_ai.response.model": "gpt-4o-2024-08-06",
            "gen_ai.provider.name": "openai",
            "gen_ai.usage.input_tokens": 100,
            "gen_ai.usage.output_tokens": 50,
        })

        exporter.export([span])

        mock_track_ai.assert_called_once()
        assert mock_track_ai.call_args.kwargs["model_name"] == "gpt-4o-2024-08-06"

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.calculate_cost", return_value=0.001)
    @patch(f"{_MOD}.track_ai_message")
    def test_fallback_to_request_model(
        self, mock_track_ai: MagicMock, mock_cost: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """When gen_ai.response.model is absent, gen_ai.request.model is used."""
        exporter = self._make_exporter()
        span = self._make_span({
            "gen_ai.operation.name": "chat",
            "gen_ai.request.model": "gpt-4o",
            "gen_ai.provider.name": "openai",
            "gen_ai.usage.input_tokens": 100,
            "gen_ai.usage.output_tokens": 50,
        })

        exporter.export([span])

        mock_track_ai.assert_called_once()
        assert mock_track_ai.call_args.kwargs["model_name"] == "gpt-4o"

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.calculate_cost", return_value=0.001)
    @patch(f"{_MOD}.track_ai_message")
    def test_cache_tokens_mapped(
        self, mock_track_ai: MagicMock, mock_cost: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """Cache read/creation tokens should be mapped when present."""
        exporter = self._make_exporter()
        span = self._make_span({
            "gen_ai.operation.name": "chat",
            "gen_ai.request.model": "claude-3-5-sonnet",
            "gen_ai.provider.name": "anthropic",
            "gen_ai.usage.input_tokens": 200,
            "gen_ai.usage.output_tokens": 100,
            "gen_ai.usage.cache_read.input_tokens": 50,
            "gen_ai.usage.cache_creation.input_tokens": 30,
        })

        exporter.export([span])

        mock_track_ai.assert_called_once()
        kwargs = mock_track_ai.call_args.kwargs
        assert kwargs["cache_read_input_tokens"] == 50
        assert kwargs["cache_creation_input_tokens"] == 30

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.calculate_cost", return_value=None)
    @patch(f"{_MOD}.track_ai_message")
    def test_total_tokens_computed(
        self, mock_track_ai: MagicMock, mock_cost: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """totalTokens = inputTokens + outputTokens."""
        exporter = self._make_exporter()
        span = self._make_span({
            "gen_ai.operation.name": "chat",
            "gen_ai.request.model": "gpt-4o",
            "gen_ai.provider.name": "openai",
            "gen_ai.usage.input_tokens": 150,
            "gen_ai.usage.output_tokens": 75,
        })

        exporter.export([span])

        mock_track_ai.assert_called_once()
        kwargs = mock_track_ai.call_args.kwargs
        assert kwargs["input_tokens"] == 150
        assert kwargs["output_tokens"] == 75
        assert kwargs["total_tokens"] == 225

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.calculate_cost", return_value=None)
    @patch(f"{_MOD}.track_tool_call")
    def test_tool_operation(
        self, mock_track_tool: MagicMock, mock_cost: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """Spans with operation='execute_tool' should call track_tool_call."""
        exporter = self._make_exporter()
        span = self._make_span({
            "gen_ai.operation.name": "execute_tool",
            "gen_ai.tool.name": "calculator",
            "gen_ai.request.model": "gpt-4o",
            "gen_ai.provider.name": "openai",
        })

        exporter.export([span])

        mock_track_tool.assert_called_once()
        kwargs = mock_track_tool.call_args.kwargs
        assert kwargs["tool_name"] == "calculator"
        assert kwargs["success"] is True
        assert kwargs["user_id"] == "test-user"

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.calculate_cost", return_value=None)
    @patch(f"{_MOD}.track_tool_call")
    def test_tool_operation_with_error(
        self, mock_track_tool: MagicMock, mock_cost: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """Tool spans with an error.type attribute should report success=False."""
        exporter = self._make_exporter()
        span = self._make_span({
            "gen_ai.operation.name": "execute_tool",
            "gen_ai.tool.name": "web_search",
            "gen_ai.request.model": "gpt-4o",
            "gen_ai.provider.name": "openai",
            "error.type": "TimeoutError",
        })

        exporter.export([span])

        mock_track_tool.assert_called_once()
        kwargs = mock_track_tool.call_args.kwargs
        assert kwargs["success"] is False
        assert kwargs["error_message"] == "TimeoutError"

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.calculate_cost", return_value=None)
    @patch(f"{_MOD}.track_embedding")
    def test_embedding_operation(
        self, mock_track_embed: MagicMock, mock_cost: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """Spans with operation='embeddings' should call track_embedding."""
        exporter = self._make_exporter()
        span = self._make_span({
            "gen_ai.operation.name": "embeddings",
            "gen_ai.request.model": "text-embedding-3-small",
            "gen_ai.provider.name": "openai",
            "gen_ai.usage.input_tokens": 50,
            "gen_ai.embeddings.dimension.count": 1536,
        })

        exporter.export([span])

        mock_track_embed.assert_called_once()
        kwargs = mock_track_embed.call_args.kwargs
        assert kwargs["model"] == "text-embedding-3-small"
        assert kwargs["provider"] == "openai"
        assert kwargs["input_tokens"] == 50
        assert kwargs["dimensions"] == 1536

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.calculate_cost", return_value=None)
    @patch(f"{_MOD}.track_ai_message")
    def test_latency_computed_from_span_timestamps(
        self, mock_track_ai: MagicMock, mock_cost: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """Latency should be computed from (end_time - start_time) in ms."""
        exporter = self._make_exporter()
        span = self._make_span(
            attrs={
                "gen_ai.operation.name": "chat",
                "gen_ai.request.model": "gpt-4o",
                "gen_ai.provider.name": "openai",
            },
            start_ns=5_000_000_000,
            end_ns=5_500_000_000,
        )

        exporter.export([span])

        mock_track_ai.assert_called_once()
        assert mock_track_ai.call_args.kwargs["latency_ms"] == 500.0

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.calculate_cost", return_value=None)
    @patch(f"{_MOD}.track_ai_message")
    def test_non_genai_span_ignored(
        self, mock_track_ai: MagicMock, mock_cost: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """Spans without gen_ai.* attributes should be silently skipped."""
        exporter = self._make_exporter()
        span = self._make_span(attrs={"http.method": "GET", "http.url": "/health"})

        result = exporter.export([span])

        mock_track_ai.assert_not_called()
        assert result == "SUCCESS"

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.track_ai_message")
    def test_bedrock_provider_passed_to_calculate_cost(
        self, mock_track_ai: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """calculate_cost should receive default_provider from the OTEL provider attribute."""
        exporter = self._make_exporter()
        span = self._make_span({
            "gen_ai.operation.name": "chat",
            "gen_ai.request.model": "anthropic.claude-sonnet-4-6",
            "gen_ai.provider.name": "bedrock",
            "gen_ai.usage.input_tokens": 100_000,
            "gen_ai.usage.output_tokens": 10_000,
        })

        with patch(f"{_MOD}.calculate_cost", return_value=0.42) as mock_cost:
            exporter.export([span])

            mock_cost.assert_called_once()
            kwargs = mock_cost.call_args.kwargs
            assert kwargs["default_provider"] == "bedrock"
            assert kwargs["model_name"] == "anthropic.claude-sonnet-4-6"

    @patch(f"{_MOD}.get_active_context", return_value=None)
    @patch(f"{_MOD}.track_ai_message")
    def test_unknown_provider_not_passed_to_calculate_cost(
        self, mock_track_ai: MagicMock, mock_ctx: MagicMock
    ) -> None:
        """When provider is 'unknown', default_provider should be None."""
        exporter = self._make_exporter()
        span = self._make_span({
            "gen_ai.operation.name": "chat",
            "gen_ai.request.model": "gpt-4o",
            "gen_ai.usage.input_tokens": 100,
            "gen_ai.usage.output_tokens": 50,
        })

        with patch(f"{_MOD}.calculate_cost", return_value=0.01) as mock_cost:
            exporter.export([span])

            mock_cost.assert_called_once()
            kwargs = mock_cost.call_args.kwargs
            assert kwargs["default_provider"] is None


class TestAmplitudeSpanExporter:
    """Tests for outbound Amplitude event -> OTEL span conversion."""

    @pytest.fixture(autouse=True)
    def _patch_otel_symbols(self) -> Any:
        """Provide mock OTEL tracer and status code."""
        self.mock_span = MagicMock()
        self.mock_tracer = MagicMock()
        self.mock_tracer.start_span.return_value = self.mock_span

        mock_trace = MagicMock()
        mock_trace.get_tracer.return_value = self.mock_tracer

        self.mock_status_code = MagicMock()
        self.mock_status_code.OK = "OK"
        self.mock_status_code.ERROR = "ERROR"

        with (
            patch(f"{_MOD}.OTEL_AVAILABLE", True),
            patch(f"{_MOD}.trace", mock_trace),
            patch(f"{_MOD}.StatusCode", self.mock_status_code),
        ):
            yield

    def _make_event(
        self,
        event_type: str = "[Agent] AI Response",
        props: dict[str, Any] | None = None,
        user_id: str = "user-1",
    ) -> MagicMock:
        """Create a mock Amplitude BaseEvent."""
        event = MagicMock()
        event.event_type = event_type
        event.event_properties = props or {}
        event.user_id = user_id
        return event

    def test_basic_conversion(self) -> None:
        """A basic Amplitude AI event should be converted to an OTEL-compatible span."""
        from amplitude_ai.integrations.opentelemetry import AmplitudeSpanExporter

        exporter = AmplitudeSpanExporter()
        event = self._make_event(props={
            "[Agent] Model Name": "gpt-4o",
            "[Agent] Provider": "openai",
            "[Agent] Session ID": "sess-123",
            "[Agent] Latency Ms": 500.0,
        })

        exporter.record_event(event)

        self.mock_tracer.start_span.assert_called_once()
        call_kwargs = self.mock_tracer.start_span.call_args.kwargs
        assert call_kwargs["name"] == "[Agent] AI Response"

        attrs = call_kwargs["attributes"]
        assert attrs["amp.[Agent] Model Name"] == "gpt-4o"
        assert attrs["amp.[Agent] Provider"] == "openai"
        assert attrs["amp.[Agent] Session ID"] == "sess-123"
        assert attrs["amp.[Agent] Latency Ms"] == 500.0

        self.mock_span.set_status.assert_called_once_with("OK")
        self.mock_span.set_attribute.assert_any_call("enduser.id", "user-1")
        self.mock_span.set_attribute.assert_any_call("amp.session_id", "sess-123")
        self.mock_span.set_attribute.assert_any_call("gen_ai.request.model", "gpt-4o")
        self.mock_span.set_attribute.assert_any_call("gen_ai.system", "openai")
        self.mock_span.set_attribute.assert_any_call("gen_ai.response.latency_ms", 500.0)
        self.mock_span.end.assert_called_once()

    def test_error_status_on_4xx(self) -> None:
        """A 4xx status_code should set the span status to ERROR."""
        from amplitude_ai.integrations.opentelemetry import AmplitudeSpanExporter

        exporter = AmplitudeSpanExporter()
        event = self._make_event(props={"[Agent] Model Name": "gpt-4o"})

        exporter.record_event(event, status_code=429, message="Rate limited")

        self.mock_span.set_status.assert_called_once_with(
            "ERROR", "Rate limited"
        )
        self.mock_span.end.assert_called_once()

    def test_error_status_on_is_error_prop(self) -> None:
        """PROP_IS_ERROR in event_properties should set ERROR status."""
        from amplitude_ai.integrations.opentelemetry import AmplitudeSpanExporter

        exporter = AmplitudeSpanExporter()
        event = self._make_event(props={
            "[Agent] Model Name": "gpt-4o",
            "[Agent] Is Error": True,
        })

        exporter.record_event(event)

        self.mock_span.set_status.assert_called_once_with(
            "ERROR", "Event delivery failed"
        )

    def test_flatten_nested_properties(self) -> None:
        """Nested dicts should be flattened with dot-separated keys."""
        from amplitude_ai.integrations.opentelemetry import AmplitudeSpanExporter

        result = AmplitudeSpanExporter._flatten_properties(
            {"top": {"nested": "value", "num": 42}},
            prefix="amp",
        )
        assert result == {"amp.top.nested": "value", "amp.top.num": 42}
