from __future__ import annotations

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

import pytest
from amplitude import Amplitude

from amplitude_ai import (
    EVENT_USER_MESSAGE,
    MockAmplitudeAI,
    PROP_MESSAGE_SOURCE,
    PROP_PARENT_AGENT_ID,
)
from amplitude_ai.core.tracking import track_user_message

Props = Dict[str, Any]


class TestAmplitudeAIMessageSource:
    """Tests for message_source via the high-level AmplitudeAI client."""

    def test_track_user_message_defaults_to_user(self) -> None:
        mock = MockAmplitudeAI()
        mock.track_user_message(
            user_id="u1", content="hello", session_id="s1",
        )

        events = mock.get_events(EVENT_USER_MESSAGE)
        assert len(events) == 1
        props: Props = events[0].event_properties or {}
        assert props.get(PROP_MESSAGE_SOURCE) == "user"

    def test_track_user_message_allows_override(self) -> None:
        mock = MockAmplitudeAI()
        mock.track_user_message(
            user_id="u1",
            content="delegated task",
            session_id="s1",
            message_source="agent",
        )

        events = mock.get_events(EVENT_USER_MESSAGE)
        assert len(events) == 1
        props: Props = events[0].event_properties or {}
        assert props.get(PROP_MESSAGE_SOURCE) == "agent"


class TestCoreTrackingMessageSource:
    """Tests for message_source via the lower-level core tracking function."""

    def test_message_source_none_omits_property(self) -> None:
        amp = MagicMock(spec=Amplitude)
        track_user_message(
            amplitude=amp,
            user_id="u1",
            message_content="hello",
            session_id="s1",
            message_source=None,
        )

        amp.track.assert_called_once()
        event = amp.track.call_args[0][0]
        props: Props = event.event_properties or {}
        assert PROP_MESSAGE_SOURCE not in props


class TestProviderMessageSource:
    """Tests for message_source set by provider wrapper's track_completion."""

    def test_no_parent_agent_sets_user(self) -> None:
        mock = MockAmplitudeAI()
        agent = mock.agent("my-agent", user_id="u1")

        with agent.session(session_id="s1") as s:
            from amplitude_ai.providers.base import BaseAIProvider

            class StubProvider(BaseAIProvider):
                def _get_provider_name(self) -> str:
                    return "stub"

                def _extract_model_name(self, **kwargs: Any) -> str:
                    return "stub-model"

                def _extract_messages(self, **kwargs: Any) -> list[dict[str, Any]]:
                    return [{"role": "user", "content": "hi"}]

                def _format_messages_for_tracking(
                    self, messages: list[dict[str, Any]],
                ) -> list[dict[str, Any]]:
                    return messages

                def _extract_response_content(self, response: Any) -> str:
                    return "hello back"

                def _extract_finish_reason(self, response: Any) -> str | None:
                    return "stop"

                def _extract_tool_calls(
                    self, response: Any,
                ) -> list[dict[str, Any]] | None:
                    return None

            provider = StubProvider(amplitude=mock.amplitude)

            fake_response = MagicMock()
            fake_response.usage = None
            fake_response.model = "stub-model"

            provider.track_completion(
                user_id="u1",
                response=fake_response,
                latency_ms=100.0,
                messages=[{"role": "user", "content": "hi"}],
            )

        user_events = mock.get_events(EVENT_USER_MESSAGE)
        assert len(user_events) >= 1
        props: Props = user_events[-1].event_properties or {}
        assert props.get(PROP_MESSAGE_SOURCE) == "user"

    def test_parent_agent_sets_agent(self) -> None:
        mock = MockAmplitudeAI()
        parent = mock.agent("orchestrator", user_id="u1")
        child = parent.child("researcher")

        with parent.session(session_id="s1") as s:
            with s.run_as(child) as cs:
                from amplitude_ai.providers.base import BaseAIProvider

                class StubProvider(BaseAIProvider):
                    def _get_provider_name(self) -> str:
                        return "stub"

                    def _extract_model_name(self, **kwargs: Any) -> str:
                        return "stub-model"

                    def _extract_messages(
                        self, **kwargs: Any,
                    ) -> list[dict[str, Any]]:
                        return [{"role": "user", "content": "sub-task"}]

                    def _format_messages_for_tracking(
                        self, messages: list[dict[str, Any]],
                    ) -> list[dict[str, Any]]:
                        return messages

                    def _extract_response_content(self, response: Any) -> str:
                        return "done"

                    def _extract_finish_reason(self, response: Any) -> str | None:
                        return "stop"

                    def _extract_tool_calls(
                        self, response: Any,
                    ) -> list[dict[str, Any]] | None:
                        return None

                provider = StubProvider(amplitude=mock.amplitude)

                fake_response = MagicMock()
                fake_response.usage = None
                fake_response.model = "stub-model"

                provider.track_completion(
                    user_id="u1",
                    response=fake_response,
                    latency_ms=50.0,
                    messages=[{"role": "user", "content": "sub-task"}],
                )

        user_events = [
            e for e in mock.get_events(EVENT_USER_MESSAGE)
            if (e.event_properties or {}).get(PROP_PARENT_AGENT_ID) is not None
        ]
        assert len(user_events) >= 1
        props: Props = user_events[0].event_properties or {}
        assert props.get(PROP_MESSAGE_SOURCE) == "agent"
