from __future__ import annotations

import asyncio
from typing import Any, Dict

import pytest

from amplitude_ai import (
    EVENT_AI_RESPONSE,
    EVENT_SESSION_END,
    EVENT_TOOL_CALL,
    EVENT_USER_MESSAGE,
    MockAmplitudeAI,
    PROP_AGENT_ID,
    PROP_PARENT_AGENT_ID,
    PROP_SESSION_ID,
    PROP_TURN_ID,
)
from amplitude_ai.context import get_active_context

Props = Dict[str, Any]


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

        captured_agent_id: str | None = None
        captured_parent_id: str | None = None

        with parent.session(session_id="s1") as s:
            with s.run_as(child) as _cs:
                ctx = get_active_context()
                captured_agent_id = ctx.agent_id if ctx else None
                captured_parent_id = ctx.parent_agent_id if ctx else None

        assert captured_agent_id == "researcher"
        assert captured_parent_id == "orchestrator"

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

        captured_session_id: str | None = None

        with parent.session(session_id="parent-session") as s:
            with s.run_as(child) as _cs:
                ctx = get_active_context()
                captured_session_id = ctx.session_id if ctx else None

        assert captured_session_id == "parent-session"

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

        captured_trace_id: str | None = None

        with parent.session(session_id="s1") as s:
            trace_id = s.new_trace()
            with s.run_as(child) as _cs:
                ctx = get_active_context()
                captured_trace_id = ctx.trace_id if ctx else None

        assert captured_trace_id == trace_id

    def test_does_not_emit_session_end(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:
                cs.track_user_message(content="sub-task")

        end_events = mock.get_events(EVENT_SESSION_END)
        assert len(end_events) == 1
        props: Props = end_events[0].event_properties or {}
        assert props.get(PROP_AGENT_ID) == "orchestrator"

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

        with parent.session(session_id="s1") as s:
            s.track_user_message(content="parent question")

            with s.run_as(child) as cs:
                cs.track_user_message(content="child sub-task")
                cs.track_ai_message(
                    content="child response",
                    model="gpt-4o",
                    provider="openai",
                    latency_ms=100,
                )
                cs.track_tool_call(tool_name="search", latency_ms=50, success=True)

            s.track_ai_message(
                content="parent response",
                model="gpt-4o",
                provider="openai",
                latency_ms=200,
            )

        user_msgs = mock.get_events(EVENT_USER_MESSAGE)
        ai_msgs = mock.get_events(EVENT_AI_RESPONSE)
        tool_calls = mock.get_events(EVENT_TOOL_CALL)

        p0: Props = user_msgs[0].event_properties or {}
        p1: Props = user_msgs[1].event_properties or {}
        assert p0.get(PROP_AGENT_ID) == "orchestrator"
        assert p1.get(PROP_AGENT_ID) == "researcher"
        assert p1.get(PROP_PARENT_AGENT_ID) == "orchestrator"

        a0: Props = ai_msgs[0].event_properties or {}
        a1: Props = ai_msgs[1].event_properties or {}
        assert a0.get(PROP_AGENT_ID) == "researcher"
        assert a1.get(PROP_AGENT_ID) == "orchestrator"

        t0: Props = tool_calls[0].event_properties or {}
        assert t0.get(PROP_AGENT_ID) == "researcher"

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

        with parent.session(session_id="shared-session") as s:
            s.track_user_message(content="q1")
            with s.run_as(child) as cs:
                cs.track_user_message(content="sub-q")
                cs.track_ai_message(
                    content="sub-a", model="gpt-4o", provider="openai", latency_ms=50,
                )
            s.track_ai_message(
                content="a1", model="gpt-4o", provider="openai", latency_ms=100,
            )

        session_events = [
            e for e in mock.events
            if (e.event_properties or {}).get(PROP_SESSION_ID) == "shared-session"
        ]
        non_end = [e for e in session_events if e.event_type != EVENT_SESSION_END]
        assert len(non_end) == 4

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

        with parent.session(session_id="s1") as s:
            s.track_user_message(content="parent msg 1")
            with s.run_as(child) as cs:
                cs.track_user_message(content="child msg")
                cs.track_ai_message(
                    content="child response",
                    model="gpt-4o",
                    provider="openai",
                    latency_ms=50,
                )
            s.track_ai_message(
                content="parent response",
                model="gpt-4o",
                provider="openai",
                latency_ms=100,
            )

        session_events = [
            e for e in mock.events
            if (e.event_properties or {}).get(PROP_SESSION_ID) == "s1"
            and e.event_type != EVENT_SESSION_END
        ]
        turn_ids = [
            (e.event_properties or {}).get(PROP_TURN_ID) for e in session_events
        ]
        for i in range(1, len(turn_ids)):
            assert turn_ids[i] > turn_ids[i - 1]  # type: ignore[operator]

    def test_nested_run_as(self) -> None:
        mock = MockAmplitudeAI()
        grandparent = mock.agent("orchestrator", user_id="u1")
        mid = grandparent.child("researcher")
        leaf = mid.child("fact-checker")

        with grandparent.session(session_id="s1") as s:
            with s.run_as(mid) as ms:
                with ms.run_as(leaf) as ls:
                    ctx = get_active_context()
                    ls.track_user_message(content="deep check")

        assert ctx is not None
        assert ctx.agent_id == "fact-checker"
        assert ctx.parent_agent_id == "researcher"
        assert ctx.session_id == "s1"

        user_msg = mock.get_events(EVENT_USER_MESSAGE)[0]
        p: Props = user_msg.event_properties or {}
        assert p.get(PROP_AGENT_ID) == "fact-checker"
        assert p.get(PROP_PARENT_AGENT_ID) == "researcher"

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

        with parent.session(session_id="s1") as s:
            before = get_active_context()
            assert before is not None
            assert before.agent_id == "orchestrator"

            with s.run_as(child) as _cs:
                assert get_active_context().agent_id == "researcher"  # type: ignore[union-attr]

            after = get_active_context()
            assert after is not None
            assert after.agent_id == "orchestrator"

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

        with parent.session(session_id="s1") as s:
            with pytest.raises(ValueError, match="child exploded"):
                with s.run_as(child) as _cs:
                    raise ValueError("child exploded")

            assert get_active_context().agent_id == "orchestrator"  # type: ignore[union-attr]

        end_events = mock.get_events(EVENT_SESSION_END)
        assert len(end_events) == 1


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

        async def _run() -> None:
            async with parent.session(session_id="s1") as s:
                async with s.arun_as(child) as cs:
                    ctx = get_active_context()
                    assert ctx is not None
                    assert ctx.agent_id == "researcher"
                    assert ctx.parent_agent_id == "orchestrator"
                    assert ctx.session_id == "s1"
                    cs.track_user_message(content="async child msg")

        asyncio.get_event_loop().run_until_complete(_run())

        user_msg = mock.get_events(EVENT_USER_MESSAGE)[0]
        p: Props = user_msg.event_properties or {}
        assert p.get(PROP_AGENT_ID) == "researcher"

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

        async def _run() -> None:
            async with parent.session(session_id="s1") as s:
                async with s.arun_as(child) as cs:
                    cs.track_user_message(content="task")

        asyncio.get_event_loop().run_until_complete(_run())

        end_events = mock.get_events(EVENT_SESSION_END)
        assert len(end_events) == 1
        p: Props = end_events[0].event_properties or {}
        assert p.get(PROP_AGENT_ID) == "orchestrator"

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

        async def _run() -> None:
            async with parent.session(session_id="s1") as s:
                with pytest.raises(ValueError, match="boom"):
                    async with s.arun_as(child) as _cs:
                        raise ValueError("boom")
                assert get_active_context().agent_id == "orchestrator"  # type: ignore[union-attr]

        asyncio.get_event_loop().run_until_complete(_run())
