"""
Unit tests for provider trace telemetry across agent versions and adapters.

Tests verify that telemetry context is correctly passed through:
- Tool generation endpoint
- LettaAgent (v1), LettaAgentV2, LettaAgentV3
- Streaming and non-streaming paths
- Different stream adapters
"""

import uuid
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from letta.schemas.llm_config import LLMConfig


@pytest.fixture
def mock_llm_config():
    """Create a mock LLM config."""
    return LLMConfig(
        model="gpt-4o-mini",
        model_endpoint_type="openai",
        model_endpoint="https://api.openai.com/v1",
        context_window=8000,
    )


class TestToolGenerationTelemetry:
    """Tests for tool generation endpoint telemetry."""

    @pytest.mark.asyncio
    async def test_generate_tool_sets_call_type(self, mock_llm_config):
        """Verify generate_tool endpoint sets call_type='tool_generation'."""
        from letta.llm_api.llm_client import LLMClient
        from letta.schemas.user import User

        mock_actor = User(
            id=f"user-{uuid.uuid4()}",
            organization_id=f"org-{uuid.uuid4()}",
            name="test_user",
        )

        captured_telemetry = {}

        def capture_telemetry(**kwargs):
            captured_telemetry.update(kwargs)

        with patch.object(LLMClient, "create") as mock_create:
            mock_client = MagicMock()
            mock_client.set_telemetry_context = capture_telemetry
            mock_client.build_request_data = MagicMock(return_value={})
            mock_client.request_async_with_telemetry = AsyncMock(return_value={})
            mock_client.convert_response_to_chat_completion = AsyncMock(
                return_value=MagicMock(
                    choices=[
                        MagicMock(
                            message=MagicMock(
                                tool_calls=[
                                    MagicMock(
                                        function=MagicMock(
                                            arguments='{"raw_source_code": "def test(): pass", "sample_args_json": "{}", "pip_requirements_json": "{}"}'
                                        )
                                    )
                                ],
                                content=None,
                            )
                        )
                    ]
                )
            )
            mock_create.return_value = mock_client

            from letta.server.rest_api.routers.v1.tools import GenerateToolInput, generate_tool_from_prompt

            mock_server = MagicMock()
            mock_server.user_manager.get_actor_or_default_async = AsyncMock(return_value=mock_actor)
            mock_server.get_llm_config_from_handle_async = AsyncMock(return_value=mock_llm_config)

            mock_headers = MagicMock()
            mock_headers.actor_id = mock_actor.id

            request = GenerateToolInput(
                prompt="Create a function that adds two numbers",
                tool_name="add_numbers",
                validation_errors=[],
            )

            with patch("letta.server.rest_api.routers.v1.tools.derive_openai_json_schema") as mock_schema:
                mock_schema.return_value = {"name": "add_numbers", "parameters": {}}
                try:
                    await generate_tool_from_prompt(request=request, server=mock_server, headers=mock_headers)
                except Exception:
                    pass

            assert captured_telemetry.get("call_type") == "tool_generation"

    @pytest.mark.asyncio
    async def test_generate_tool_has_no_agent_context(self, mock_llm_config):
        """Verify generate_tool doesn't have agent_id since it's not agent-bound."""
        from letta.llm_api.llm_client import LLMClient
        from letta.schemas.user import User

        mock_actor = User(
            id=f"user-{uuid.uuid4()}",
            organization_id=f"org-{uuid.uuid4()}",
            name="test_user",
        )

        captured_telemetry = {}

        def capture_telemetry(**kwargs):
            captured_telemetry.update(kwargs)

        with patch.object(LLMClient, "create") as mock_create:
            mock_client = MagicMock()
            mock_client.set_telemetry_context = capture_telemetry
            mock_client.build_request_data = MagicMock(return_value={})
            mock_client.request_async_with_telemetry = AsyncMock(return_value={})
            mock_client.convert_response_to_chat_completion = AsyncMock(
                return_value=MagicMock(
                    choices=[
                        MagicMock(
                            message=MagicMock(
                                tool_calls=[
                                    MagicMock(
                                        function=MagicMock(
                                            arguments='{"raw_source_code": "def test(): pass", "sample_args_json": "{}", "pip_requirements_json": "{}"}'
                                        )
                                    )
                                ],
                                content=None,
                            )
                        )
                    ]
                )
            )
            mock_create.return_value = mock_client

            from letta.server.rest_api.routers.v1.tools import GenerateToolInput, generate_tool_from_prompt

            mock_server = MagicMock()
            mock_server.user_manager.get_actor_or_default_async = AsyncMock(return_value=mock_actor)
            mock_server.get_llm_config_from_handle_async = AsyncMock(return_value=mock_llm_config)

            mock_headers = MagicMock()
            mock_headers.actor_id = mock_actor.id

            request = GenerateToolInput(
                prompt="Create a function",
                tool_name="test_func",
                validation_errors=[],
            )

            with patch("letta.server.rest_api.routers.v1.tools.derive_openai_json_schema") as mock_schema:
                mock_schema.return_value = {"name": "test_func", "parameters": {}}
                try:
                    await generate_tool_from_prompt(request=request, server=mock_server, headers=mock_headers)
                except Exception:
                    pass

            assert captured_telemetry.get("agent_id") is None
            assert captured_telemetry.get("step_id") is None
            assert captured_telemetry.get("run_id") is None


class TestLLMClientTelemetryContext:
    """Tests for LLMClient telemetry context methods."""

    def test_llm_client_has_set_telemetry_context_method(self):
        """Verify LLMClient exposes set_telemetry_context."""
        from letta.llm_api.llm_client import LLMClient

        client = LLMClient.create(provider_type="openai", put_inner_thoughts_first=True)
        assert hasattr(client, "set_telemetry_context")
        assert callable(client.set_telemetry_context)

    def test_llm_client_set_telemetry_context_accepts_all_fields(self):
        """Verify set_telemetry_context accepts all telemetry fields."""
        from letta.llm_api.llm_client import LLMClient

        client = LLMClient.create(provider_type="openai", put_inner_thoughts_first=True)

        client.set_telemetry_context(
            agent_id=f"agent-{uuid.uuid4()}",
            agent_tags=["tag1", "tag2"],
            run_id=f"run-{uuid.uuid4()}",
            step_id=f"step-{uuid.uuid4()}",
            call_type="summarization",
        )


class TestAdapterTelemetryAttributes:
    """Tests for adapter telemetry attribute support."""

    def test_base_adapter_has_telemetry_attributes(self, mock_llm_config):
        """Verify base LettaLLMAdapter has telemetry attributes."""
        from letta.adapters.letta_llm_adapter import LettaLLMAdapter
        from letta.llm_api.llm_client import LLMClient
        from letta.schemas.enums import LLMCallType

        mock_client = LLMClient.create(provider_type="openai", put_inner_thoughts_first=True)

        agent_id = f"agent-{uuid.uuid4()}"
        agent_tags = ["test-tag"]
        run_id = f"run-{uuid.uuid4()}"

        class TestAdapter(LettaLLMAdapter):
            async def invoke_llm(self, *args, **kwargs):
                pass

        adapter = TestAdapter(
            llm_client=mock_client,
            llm_config=mock_llm_config,
            call_type=LLMCallType.agent_step,
            agent_id=agent_id,
            agent_tags=agent_tags,
            run_id=run_id,
        )

        assert adapter.agent_id == agent_id
        assert adapter.agent_tags == agent_tags
        assert adapter.run_id == run_id
        assert adapter.call_type == LLMCallType.agent_step

    def test_request_adapter_inherits_telemetry_attributes(self, mock_llm_config):
        """Verify LettaLLMRequestAdapter inherits telemetry attributes."""
        from letta.adapters.letta_llm_request_adapter import LettaLLMRequestAdapter
        from letta.llm_api.llm_client import LLMClient
        from letta.schemas.enums import LLMCallType

        mock_client = LLMClient.create(provider_type="openai", put_inner_thoughts_first=True)

        agent_id = f"agent-{uuid.uuid4()}"
        agent_tags = ["request-tag"]
        run_id = f"run-{uuid.uuid4()}"

        adapter = LettaLLMRequestAdapter(
            llm_client=mock_client,
            llm_config=mock_llm_config,
            call_type=LLMCallType.agent_step,
            agent_id=agent_id,
            agent_tags=agent_tags,
            run_id=run_id,
        )

        assert adapter.agent_id == agent_id
        assert adapter.agent_tags == agent_tags
        assert adapter.run_id == run_id

    def test_stream_adapter_inherits_telemetry_attributes(self, mock_llm_config):
        """Verify LettaLLMStreamAdapter inherits telemetry attributes."""
        from letta.adapters.letta_llm_stream_adapter import LettaLLMStreamAdapter
        from letta.llm_api.llm_client import LLMClient
        from letta.schemas.enums import LLMCallType

        mock_client = LLMClient.create(provider_type="openai", put_inner_thoughts_first=True)

        agent_id = f"agent-{uuid.uuid4()}"
        agent_tags = ["stream-tag"]
        run_id = f"run-{uuid.uuid4()}"

        adapter = LettaLLMStreamAdapter(
            llm_client=mock_client,
            llm_config=mock_llm_config,
            call_type=LLMCallType.agent_step,
            agent_id=agent_id,
            agent_tags=agent_tags,
            run_id=run_id,
        )

        assert adapter.agent_id == agent_id
        assert adapter.agent_tags == agent_tags
        assert adapter.run_id == run_id

    def test_request_and_stream_adapters_have_consistent_interface(self, mock_llm_config):
        """Verify both adapter types have the same telemetry interface."""
        from letta.adapters.letta_llm_request_adapter import LettaLLMRequestAdapter
        from letta.adapters.letta_llm_stream_adapter import LettaLLMStreamAdapter
        from letta.llm_api.llm_client import LLMClient
        from letta.schemas.enums import LLMCallType

        mock_client = LLMClient.create(provider_type="openai", put_inner_thoughts_first=True)

        request_adapter = LettaLLMRequestAdapter(llm_client=mock_client, llm_config=mock_llm_config, call_type=LLMCallType.agent_step)
        stream_adapter = LettaLLMStreamAdapter(llm_client=mock_client, llm_config=mock_llm_config, call_type=LLMCallType.agent_step)

        for attr in ["agent_id", "agent_tags", "run_id", "call_type"]:
            assert hasattr(request_adapter, attr), f"LettaLLMRequestAdapter missing {attr}"
            assert hasattr(stream_adapter, attr), f"LettaLLMStreamAdapter missing {attr}"


class TestSummarizerTelemetry:
    """Tests for Summarizer class telemetry context."""

    def test_summarizer_stores_telemetry_context(self):
        """Verify Summarizer stores telemetry context from constructor."""
        from letta.schemas.user import User
        from letta.services.summarizer.enums import SummarizationMode
        from letta.services.summarizer.summarizer import Summarizer

        mock_actor = User(
            id=f"user-{uuid.uuid4()}",
            organization_id=f"org-{uuid.uuid4()}",
            name="test_user",
        )

        agent_id = f"agent-{uuid.uuid4()}"
        run_id = f"run-{uuid.uuid4()}"
        step_id = f"step-{uuid.uuid4()}"

        summarizer = Summarizer(
            mode=SummarizationMode.PARTIAL_EVICT_MESSAGE_BUFFER,
            summarizer_agent=None,
            message_buffer_limit=100,
            message_buffer_min=10,
            partial_evict_summarizer_percentage=0.5,
            agent_manager=MagicMock(),
            message_manager=MagicMock(),
            actor=mock_actor,
            agent_id=agent_id,
            run_id=run_id,
            step_id=step_id,
        )

        assert summarizer.agent_id == agent_id
        assert summarizer.run_id == run_id
        assert summarizer.step_id == step_id

    @pytest.mark.asyncio
    async def test_summarize_method_accepts_runtime_telemetry(self):
        """Verify summarize() method accepts runtime run_id/step_id."""
        from letta.schemas.enums import MessageRole
        from letta.schemas.message import Message
        from letta.schemas.user import User
        from letta.services.summarizer.enums import SummarizationMode
        from letta.services.summarizer.summarizer import Summarizer

        mock_actor = User(
            id=f"user-{uuid.uuid4()}",
            organization_id=f"org-{uuid.uuid4()}",
            name="test_user",
        )

        agent_id = f"agent-{uuid.uuid4()}"
        mock_messages = [
            Message(
                id=f"message-{uuid.uuid4()}",
                role=MessageRole.user,
                content=[{"type": "text", "text": "Hello"}],
                agent_id=agent_id,
            )
        ]

        summarizer = Summarizer(
            mode=SummarizationMode.PARTIAL_EVICT_MESSAGE_BUFFER,
            summarizer_agent=None,
            message_buffer_limit=100,
            message_buffer_min=10,
            partial_evict_summarizer_percentage=0.5,
            agent_manager=MagicMock(),
            message_manager=MagicMock(),
            actor=mock_actor,
            agent_id=agent_id,
        )

        run_id = f"run-{uuid.uuid4()}"
        step_id = f"step-{uuid.uuid4()}"

        result = await summarizer.summarize(
            in_context_messages=mock_messages,
            new_letta_messages=[],
            force=False,
            run_id=run_id,
            step_id=step_id,
        )

        assert result is not None


class TestAgentAdapterInstantiation:
    """Tests verifying agents instantiate adapters with telemetry context."""

    def test_agent_v2_creates_summarizer_with_agent_id(self, mock_llm_config):
        """Verify LettaAgentV2 creates Summarizer with correct agent_id."""
        from letta.agents.letta_agent_v2 import LettaAgentV2
        from letta.schemas.agent import AgentState, AgentType
        from letta.schemas.embedding_config import EmbeddingConfig
        from letta.schemas.memory import Memory
        from letta.schemas.user import User

        mock_actor = User(
            id=f"user-{uuid.uuid4()}",
            organization_id=f"org-{uuid.uuid4()}",
            name="test_user",
        )

        agent_id = f"agent-{uuid.uuid4()}"
        agent_state = AgentState(
            id=agent_id,
            name="test_agent",
            agent_type=AgentType.letta_v1_agent,
            llm_config=mock_llm_config,
            embedding_config=EmbeddingConfig.default_config(provider="openai"),
            tags=["test"],
            memory=Memory(blocks=[]),
            system="You are a helpful assistant.",
            tools=[],
            sources=[],
            blocks=[],
        )

        agent = LettaAgentV2(agent_state=agent_state, actor=mock_actor)

        assert agent.summarizer.agent_id == agent_id