"""Tests for agent handoff trace separation."""

import asyncio
import pytest
from uuid import uuid4
from unittest.mock import MagicMock, patch

from lumenova_beacon import BeaconClient
from lumenova_beacon.tracing.span import Span, SpanLink
from lumenova_beacon.types import SpanType, StatusCode


@pytest.fixture
def client():
    """Create a BeaconClient with OTEL disabled for unit tests."""
    return BeaconClient(
        endpoint="https://api.test.com",
        api_key="test-key",
        auto_instrument_opentelemetry=False,
    )


class TestSpanLink:
    """Tests for the SpanLink dataclass."""

    def test_span_link_creation(self):
        link = SpanLink(trace_id="aabb", span_id="ccdd")
        assert link.trace_id == "aabb"
        assert link.span_id == "ccdd"
        assert link.attributes == {}

    def test_span_link_with_attributes(self):
        link = SpanLink(
            trace_id="aabb",
            span_id="ccdd",
            attributes={"link.type": "child_trace"},
        )
        assert link.attributes == {"link.type": "child_trace"}


class TestSpanAddLink:
    """Tests for Span.add_link()."""

    def test_add_link(self):
        span = Span(name="test")
        span.add_link("trace-aaa", "span-bbb")
        assert len(span.links) == 1
        assert span.links[0].trace_id == "trace-aaa"
        assert span.links[0].span_id == "span-bbb"
        assert span.links[0].attributes == {}

    def test_add_link_with_attributes(self):
        span = Span(name="test")
        span.add_link("trace-aaa", "span-bbb", {"link.type": "child_trace"})
        assert span.links[0].attributes == {"link.type": "child_trace"}

    def test_add_link_returns_self(self):
        span = Span(name="test")
        result = span.add_link("trace-aaa", "span-bbb")
        assert result is span

    def test_add_multiple_links(self):
        span = Span(name="test")
        span.add_link("trace-1", "span-1")
        span.add_link("trace-2", "span-2")
        assert len(span.links) == 2

    def test_links_empty_by_default(self):
        span = Span(name="test")
        assert span.links == []


class TestSpanToDict:
    """Tests for Span.to_dict() with links."""

    def test_to_dict_without_links(self):
        span = Span(name="test")
        span.start()
        span.end()
        d = span.to_dict()
        assert "links" not in d

    def test_to_dict_with_links(self):
        span = Span(name="test")
        span.add_link("trace-aaa", "span-bbb", {"link.type": "child_trace"})
        span.start()
        span.end()
        d = span.to_dict()
        assert "links" in d
        assert len(d["links"]) == 1
        assert d["links"][0] == {
            "context": {"trace_id": "trace-aaa", "span_id": "span-bbb"},
            "attributes": {"link.type": "child_trace"},
        }


class TestSpanToOtelSpanLinks:
    """Tests for Span.to_otel_span() with links."""

    def test_otel_span_empty_links(self):
        span = Span(
            name="test",
            trace_id="a" * 32,
            span_id="b" * 16,
        )
        span.start()
        span.end()
        otel_span = span.to_otel_span()
        assert otel_span.links == ()

    def test_otel_span_with_links(self):
        span = Span(
            name="test",
            trace_id="a" * 32,
            span_id="b" * 16,
        )
        span.add_link("c" * 32, "d" * 16, {"link.type": "child_trace"})
        span.start()
        span.end()
        otel_span = span.to_otel_span()
        assert len(otel_span.links) == 1

        link = otel_span.links[0]
        assert link.context.trace_id == int("c" * 32, 16)
        assert link.context.span_id == int("d" * 16, 16)
        assert link.context.is_remote is True
        assert link.attributes == {"link.type": "child_trace"}


class TestHandlerHandoffContext:
    """Tests for BeaconLangGraphHandler.handoff_context() with ContextVar-based state isolation."""

    @pytest.fixture
    def handler(self):
        """Create a BeaconLangGraphHandler for testing."""
        from lumenova_beacon.tracing.integrations.langchain import BeaconLangGraphHandler
        return BeaconLangGraphHandler(agent_name="Orchestrator")

    def _get_handoff_state(self):
        """Get the active handoff state from ContextVar."""
        from lumenova_beacon.tracing.integrations.langchain import _active_handoff_state
        return _active_handoff_state.get()

    def test_handoff_span_in_parent_trace(self, handler):
        """HANDOFF span should use the handler's current trace_id."""
        handler._current_trace_id = "parent_trace_aaa"
        handler._root_span_id = "root_span_bbb"

        with handler.handoff_context("Math Agent") as ctx:
            assert ctx.handoff_span.trace_id == "parent_trace_aaa"
            assert ctx.handoff_span.span_type == SpanType.HANDOFF
            assert ctx.handoff_span.name == "handoff:Math Agent"

    def test_handoff_span_parented_under_active_tool(self, handler):
        """HANDOFF span parent_id should be the active tool span."""
        handler._current_trace_id = "parent_trace"
        tool_span = Span(name="handoff_to_math_agent", trace_id="parent_trace")
        handler._runs = {"tool-run-1": {"span": tool_span}}
        handler._active_tool_run_ids = ["tool-run-1"]

        with handler.handoff_context("Math Agent") as ctx:
            assert ctx.handoff_span.parent_id == tool_span.span_id

    def test_handoff_span_parented_via_config(self, handler):
        """With config, HANDOFF span should parent under the correct tool even with multiple active tools."""
        handler._current_trace_id = "parent_trace"
        tool_a_span = Span(name="tool_a", trace_id="parent_trace")
        tool_b_span = Span(name="tool_b", trace_id="parent_trace")
        handler._runs = {
            "tool-run-a": {"span": tool_a_span},
            "tool-run-b": {"span": tool_b_span},
        }
        handler._active_tool_run_ids = ["tool-run-a", "tool-run-b"]

        # Simulate config with callback manager whose parent_run_id is tool A
        mock_callbacks = MagicMock()
        mock_callbacks.parent_run_id = "tool-run-a"
        config = {"callbacks": mock_callbacks}

        with handler.handoff_context("Math Agent", config=config) as ctx:
            assert ctx.handoff_span.parent_id == tool_a_span.span_id

    def test_state_isolated_for_child(self, handler):
        """After entering, ContextVar state should be set for child trace (handler base untouched)."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "parent_root"
        handler._runs = {"run-1": {"span": Span(name="test")}}
        handler._active_tool_run_ids = ["run-1"]

        with handler.handoff_context("Math Agent") as ctx:
            hs = self._get_handoff_state()
            assert hs is not None
            assert hs.trace_id == ctx.trace_id
            assert hs.root_span_id == ctx.root_span_id
            assert hs.runs == {}
            assert hs.active_tool_run_ids == []
            # Handler base state is NOT mutated
            assert handler._current_trace_id == "parent_trace"
            assert handler._root_span_id == "parent_root"

    def test_backlink_params_set_for_child(self, handler):
        """Handoff state's back-link params should point to parent handoff span."""
        handler._current_trace_id = "parent_trace"
        handler._parent_handoff_trace_id = None
        handler._parent_handoff_span_id = None

        with handler.handoff_context("Math Agent") as ctx:
            hs = self._get_handoff_state()
            assert hs.parent_handoff_trace_id == "parent_trace"
            assert hs.parent_handoff_span_id == ctx.handoff_span.span_id

    def test_restores_state_on_exit(self, handler):
        """Handler base state should remain unchanged, ContextVar should be cleared."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "parent_root"
        original_runs = {"run-1": {"span": Span(name="test")}}
        handler._runs = original_runs
        handler._active_tool_run_ids = ["run-1"]
        handler._agent_name = "Orchestrator"

        with handler.handoff_context("Math Agent"):
            pass

        assert handler._current_trace_id == "parent_trace"
        assert handler._root_span_id == "parent_root"
        assert handler._runs is original_runs
        assert handler._active_tool_run_ids == ["run-1"]
        assert handler._agent_name == "Orchestrator"
        assert self._get_handoff_state() is None

    def test_restores_state_on_exception(self, handler):
        """ContextVar should be cleared even on exception, handler base untouched."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "parent_root"

        try:
            with handler.handoff_context("Math Agent"):
                raise ValueError("child failed")
        except ValueError:
            pass

        assert handler._current_trace_id == "parent_trace"
        assert handler._root_span_id == "parent_root"
        assert self._get_handoff_state() is None

    def test_handoff_span_has_child_trace_link(self, handler):
        """HANDOFF span should have OTEL link to the child root span."""
        handler._current_trace_id = "parent_trace"

        with handler.handoff_context("Math Agent") as ctx:
            assert len(ctx.handoff_span.links) == 1
            link = ctx.handoff_span.links[0]
            assert link.trace_id == ctx.trace_id
            assert link.span_id == ctx.root_span_id
            assert link.attributes == {"link.type": "child_trace"}

    def test_handoff_span_attributes(self, handler):
        """HANDOFF span should carry target info in attributes."""
        handler._current_trace_id = "parent_trace"

        with handler.handoff_context("Math Agent") as ctx:
            attrs = ctx.handoff_span.attributes
            assert attrs["handoff.target_agent"] == "Math Agent"
            assert attrs["handoff.target_trace_id"] == ctx.trace_id
            assert attrs["handoff.target_root_span_id"] == ctx.root_span_id

    def test_handoff_span_exported(self, handler):
        """HANDOFF span should be exported on exit."""
        handler._current_trace_id = "parent_trace"

        with patch.object(handler.client, "export_span") as mock_export:
            with handler.handoff_context("Math Agent") as ctx:
                pass
            mock_export.assert_called_once_with(ctx.handoff_span)

    def test_handoff_span_records_exception(self, handler):
        """HANDOFF span should record child exceptions."""
        handler._current_trace_id = "parent_trace"

        try:
            with handler.handoff_context("Math Agent") as ctx:
                raise RuntimeError("child crashed")
        except RuntimeError:
            pass

        assert ctx.handoff_span.status_code == StatusCode.ERROR

    def test_agent_name_set_to_target(self, handler):
        """target_agent should set agent_name in handoff state during handoff."""
        handler._current_trace_id = "parent_trace"
        handler._agent_name = "Orchestrator"

        with handler.handoff_context("Math Agent"):
            hs = self._get_handoff_state()
            assert hs.agent_name == "Math Agent"
            # Handler base agent_name unchanged
            assert handler._agent_name == "Orchestrator"

        assert handler._agent_name == "Orchestrator"

    def test_sequential_handoffs(self, handler):
        """Multiple sequential handoffs should each work independently."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "parent_root"

        with handler.handoff_context("Agent A") as ctx_a:
            trace_a = ctx_a.trace_id

        with handler.handoff_context("Agent B") as ctx_b:
            trace_b = ctx_b.trace_id

        assert trace_a != trace_b
        assert handler._current_trace_id == "parent_trace"

    @pytest.mark.asyncio
    async def test_async_handoff_context(self, handler):
        """Async handoff_context should work identically."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "parent_root"

        async with handler.handoff_context("Math Agent") as ctx:
            hs = self._get_handoff_state()
            assert hs.trace_id == ctx.trace_id
            assert ctx.handoff_span.trace_id == "parent_trace"
            # Handler base untouched
            assert handler._current_trace_id == "parent_trace"

        assert handler._current_trace_id == "parent_trace"
        assert self._get_handoff_state() is None

    @pytest.mark.asyncio
    async def test_parallel_handoffs_isolated(self, handler):
        """Two concurrent handoffs should not interfere with each other."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "parent_root"
        handler._agent_name = "Orchestrator"

        traces_seen: dict[str, str] = {}
        agent_names_seen: dict[str, str] = {}

        async def do_handoff(name: str):
            async with handler.handoff_context(name) as ctx:
                hs = self._get_handoff_state()
                traces_seen[name] = hs.trace_id
                agent_names_seen[name] = hs.agent_name
                assert hs.trace_id == ctx.trace_id
                assert hs.agent_name == name
                await asyncio.sleep(0.01)  # allow interleaving
                # Still isolated after yield
                hs2 = self._get_handoff_state()
                assert hs2.trace_id == ctx.trace_id
                assert hs2.agent_name == name

        await asyncio.gather(do_handoff("Agent A"), do_handoff("Agent B"))

        # Each got a unique trace
        assert traces_seen["Agent A"] != traces_seen["Agent B"]
        assert agent_names_seen["Agent A"] == "Agent A"
        assert agent_names_seen["Agent B"] == "Agent B"
        # Parent state untouched
        assert handler._current_trace_id == "parent_trace"
        assert handler._agent_name == "Orchestrator"
        assert self._get_handoff_state() is None

    @pytest.mark.asyncio
    async def test_parallel_handoffs_independent_runs(self, handler):
        """Parallel handoffs should have completely independent _runs dicts."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "parent_root"
        parent_runs = {"parent-run": {"span": Span(name="parent")}}
        handler._runs = parent_runs

        runs_during: dict[str, dict] = {}

        async def do_handoff(name: str):
            async with handler.handoff_context(name):
                hs = self._get_handoff_state()
                # Each handoff starts with empty runs
                assert hs.runs == {}
                # Simulate adding a run
                hs.runs[f"run-{name}"] = {"span": Span(name=name)}
                await asyncio.sleep(0.01)
                runs_during[name] = dict(hs.runs)
                # Should only see own run
                assert f"run-{name}" in hs.runs
                other = "Agent B" if name == "Agent A" else "Agent A"
                assert f"run-{other}" not in hs.runs

        await asyncio.gather(do_handoff("Agent A"), do_handoff("Agent B"))

        # Parent runs untouched
        assert handler._runs is parent_runs
        assert "parent-run" in handler._runs

    @pytest.mark.asyncio
    async def test_parallel_handoff_spans_in_parent_trace(self, handler):
        """Both parallel handoff spans should be in the parent trace."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "parent_root"

        handoff_spans: list[Span] = []

        async def do_handoff(name: str):
            async with handler.handoff_context(name) as ctx:
                handoff_spans.append(ctx.handoff_span)
                await asyncio.sleep(0.01)

        await asyncio.gather(do_handoff("Agent A"), do_handoff("Agent B"))

        assert len(handoff_spans) == 2
        for hs in handoff_spans:
            assert hs.trace_id == "parent_trace"
            assert hs.span_type == SpanType.HANDOFF

    @pytest.mark.asyncio
    async def test_parallel_handoff_spans_claim_different_tools(self, handler):
        """Each parallel handoff should parent under a different tool span (claim-based)."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "parent_root"

        # Set up a tools chain node with three child tool spans
        tools_node_span = Span(name="tools", trace_id="parent_trace", span_id="tools-span")
        tool_a_span = Span(name="handoff_to_math", trace_id="parent_trace", parent_id="tools-span")
        tool_b_span = Span(name="handoff_to_math", trace_id="parent_trace", parent_id="tools-span")
        tool_c_span = Span(name="handoff_to_math", trace_id="parent_trace", parent_id="tools-span")
        handler._runs = {
            "tools-node": {"span": tools_node_span},
            "tool-a": {"span": tool_a_span},
            "tool-b": {"span": tool_b_span},
            "tool-c": {"span": tool_c_span},
        }
        handler._active_tool_run_ids = ["tool-a", "tool-b", "tool-c"]

        # Config's parent_run_id points to tools chain node (not individual tools)
        mock_callbacks = MagicMock()
        mock_callbacks.parent_run_id = "tools-node"

        handoff_parents: list[str] = []

        async def do_handoff(name: str):
            config = {"callbacks": mock_callbacks}
            async with handler.handoff_context(name, config=config) as ctx:
                handoff_parents.append(ctx.handoff_span.parent_id)
                await asyncio.sleep(0.01)

        await asyncio.gather(
            do_handoff("Math A"), do_handoff("Math B"), do_handoff("Math C")
        )

        # Each handoff should have claimed a different tool span
        assert len(handoff_parents) == 3
        tool_span_ids = {tool_a_span.span_id, tool_b_span.span_id, tool_c_span.span_id}
        assert set(handoff_parents) == tool_span_ids


class TestHandoffInterruptHandling:
    """Tests for interrupt handling in HandoffContext."""

    @pytest.fixture
    def handler(self):
        from lumenova_beacon.tracing.integrations.langchain import BeaconLangGraphHandler
        return BeaconLangGraphHandler(agent_name="Orchestrator")

    def _get_handoff_state(self):
        from lumenova_beacon.tracing.integrations.langchain import _active_handoff_state
        return _active_handoff_state.get()

    def test_handoff_span_status_ok_on_graph_interrupt(self, handler):
        """HANDOFF span should have status OK (not ERROR) when GraphInterrupt occurs."""
        handler._current_trace_id = "parent_trace"

        class GraphInterrupt(Exception):
            pass

        try:
            with handler.handoff_context("Hangman") as ctx:
                raise GraphInterrupt("Enter a letter:")
        except GraphInterrupt:
            pass

        assert ctx.handoff_span.status_code == StatusCode.OK
        assert ctx.handoff_span.attributes.get("langgraph.interrupt") is True

    def test_handoff_span_interrupt_with_value_and_id(self, handler):
        """HANDOFF span should capture interrupt value and ID attributes."""
        handler._current_trace_id = "parent_trace"

        class Interrupt(Exception):
            def __init__(self, value):
                self.value = value
                self.id = "checkpoint-123"
                super().__init__(value)

        try:
            with handler.handoff_context("Hangman") as ctx:
                raise Interrupt("Enter a letter:")
        except Interrupt:
            pass

        assert ctx.handoff_span.status_code == StatusCode.OK
        assert ctx.handoff_span.attributes.get("langgraph.interrupt") is True
        assert ctx.handoff_span.attributes.get("langgraph.interrupt.value") == "Enter a letter:"
        assert ctx.handoff_span.attributes.get("langgraph.interrupt.id") == "checkpoint-123"
        assert "Interrupted:" in (ctx.handoff_span.status_description or "")

    def test_handoff_real_error_still_records_exception(self, handler):
        """Non-interrupt exceptions should still be recorded as errors."""
        handler._current_trace_id = "parent_trace"

        try:
            with handler.handoff_context("Math Agent") as ctx:
                raise RuntimeError("child crashed")
        except RuntimeError:
            pass

        assert ctx.handoff_span.status_code == StatusCode.ERROR

    def test_handoff_cancelled_still_ok(self, handler):
        """CancelledError should still result in OK status."""
        handler._current_trace_id = "parent_trace"
        import asyncio

        try:
            with handler.handoff_context("Math Agent") as ctx:
                raise asyncio.CancelledError()
        except asyncio.CancelledError:
            pass

        assert ctx.handoff_span.status_code == StatusCode.OK
        assert ctx.handoff_span.attributes.get("langgraph.cancelled") is True


class TestHandoffTraceContinuation:
    """Tests for trace continuation (resume) support in HandoffContext."""

    @pytest.fixture
    def handler(self):
        from lumenova_beacon.tracing.integrations.langchain import BeaconLangGraphHandler
        return BeaconLangGraphHandler(agent_name="Orchestrator")

    def _get_handoff_state(self):
        from lumenova_beacon.tracing.integrations.langchain import _active_handoff_state
        return _active_handoff_state.get()

    def test_fresh_handoff_generates_ids(self, handler):
        """Fresh handoff should generate new trace IDs and not be a resume."""
        handler._current_trace_id = "parent_trace"

        with handler.handoff_context("Hangman") as ctx:
            assert ctx.trace_id is not None
            assert len(ctx.trace_id) == 32
            assert len(ctx.root_span_id) == 16
            hs = self._get_handoff_state()
            assert hs.is_resume is False
            assert ctx.handoff_span.attributes.get("handoff.is_resume") is None

    def test_auto_resume_after_interrupt(self, handler):
        """After an interrupted handoff, the next handoff to the same agent
        should automatically reuse the child trace IDs."""
        handler._current_trace_id = "parent_trace"

        # First handoff — interrupted
        class Interrupt(Exception):
            def __init__(self, value):
                self.value = value
                self.id = "ckpt-1"
                super().__init__(value)

        try:
            with handler.handoff_context("Hangman") as ctx1:
                first_trace_id = ctx1.trace_id
                first_root_span_id = ctx1.root_span_id
                raise Interrupt("Enter a letter:")
        except Interrupt:
            pass

        # Second handoff — should auto-resume with same child trace IDs
        with handler.handoff_context("Hangman") as ctx2:
            assert ctx2.trace_id == first_trace_id
            assert ctx2.root_span_id == first_root_span_id
            assert ctx2._is_resume is True
            assert ctx2.handoff_span.attributes.get("handoff.is_resume") is True


    def test_auto_resume_cleans_up_on_normal_completion(self, handler):
        """After a normal handoff completion, saved resume state is cleaned up."""
        handler._current_trace_id = "parent_trace"

        # First handoff — interrupted
        class Interrupt(Exception):
            def __init__(self, value):
                self.value = value
                self.id = "ckpt-1"
                super().__init__(value)

        try:
            with handler.handoff_context("Hangman") as ctx1:
                first_trace_id = ctx1.trace_id
                raise Interrupt("Enter a letter:")
        except Interrupt:
            pass

        # Second handoff — completes normally
        with handler.handoff_context("Hangman") as ctx2:
            assert ctx2.trace_id == first_trace_id  # resumed
            pass  # no exception = normal completion

        # Third handoff — should get fresh IDs (resume state was cleaned up)
        with handler.handoff_context("Hangman") as ctx3:
            assert ctx3.trace_id != first_trace_id
            assert ctx3._is_resume is False


    def test_auto_resume_scoped_to_agent_name(self, handler):
        """Auto-resume should only match the same target agent name."""
        handler._current_trace_id = "parent_trace"

        class Interrupt(Exception):
            def __init__(self, value):
                self.value = value
                self.id = "ckpt-1"
                super().__init__(value)

        # Interrupt handoff to "Hangman"
        try:
            with handler.handoff_context("Hangman") as ctx1:
                hangman_trace = ctx1.trace_id
                raise Interrupt("Enter a letter:")
        except Interrupt:
            pass

        # Handoff to "Math Agent" — should NOT resume Hangman's trace
        with handler.handoff_context("Math Agent") as ctx2:
            assert ctx2.trace_id != hangman_trace
            assert ctx2._is_resume is False


    def test_resume_skips_root_span_in_handoff(self, handler):
        """On resume inside a handoff, the child graph's root chain should be
        skipped and its children reparented under the original root span."""
        from uuid import uuid4

        handler._current_trace_id = "parent_trace"

        class Interrupt(Exception):
            def __init__(self, value):
                self.value = value
                self.id = "ckpt-1"
                super().__init__(value)

        # First handoff — simulate on_chain_start to create root span, then interrupt
        try:
            with handler.handoff_context("Hangman") as ctx1:
                root_span_id = ctx1.root_span_id
                # Simulate what LangGraph does: fire on_chain_start for root chain
                # with parent_run_id from orchestrator's tool
                orchestrator_tool_run = uuid4()
                child_root_run = uuid4()
                handler.on_chain_start(
                    {"name": "Hangman Game"},
                    {"messages": []},
                    run_id=child_root_run,
                    parent_run_id=orchestrator_tool_run,
                )
                hs = self._get_handoff_state()
                # Root span should have been created with the pre-generated root_span_id
                assert str(child_root_run) in hs.runs
                assert hs.runs[str(child_root_run)]["span"].span_id == root_span_id
                # End it before interrupt
                handler.on_chain_end({"messages": []}, run_id=child_root_run)
                raise Interrupt("Enter a letter:")
        except Interrupt:
            pass

        # Second handoff (resume) — root chain should be skipped
        with handler.handoff_context("Hangman") as ctx2:
            assert ctx2._is_resume is True
            assert ctx2.root_span_id == root_span_id

            resume_root_run = uuid4()
            handler.on_chain_start(
                {"name": "Hangman Game"},
                {"messages": []},
                run_id=resume_root_run,
                parent_run_id=orchestrator_tool_run,
            )
            hs = self._get_handoff_state()
            # Root should have been SKIPPED (in skipped_resume_roots, not in runs)
            assert str(resume_root_run) in hs.skipped_resume_roots
            assert str(resume_root_run) not in hs.runs

    def test_resumed_nodes_get_chain_type_not_agent(self, handler):
        """After resume, only the root span should be AGENT type.
        Subsequent graph-step nodes must remain CHAIN type."""
        from lumenova_beacon.tracing.span import SpanType
        from uuid import uuid4

        handler._current_trace_id = "parent_trace"

        class Interrupt(Exception):
            def __init__(self, value):
                self.value = value
                self.id = "ckpt-1"
                super().__init__(value)

        orchestrator_tool_run = uuid4()

        # First handoff — create root span, one node, then interrupt
        try:
            with handler.handoff_context("Hangman") as ctx1:
                root_run = uuid4()
                handler.on_chain_start(
                    {"name": "Hangman Game"},
                    {"messages": []},
                    run_id=root_run,
                    parent_run_id=orchestrator_tool_run,
                )
                hs = self._get_handoff_state()
                root_span = hs.runs[str(root_run)]["span"]
                assert root_span.span_type == SpanType.AGENT

                # Node A
                node_a = uuid4()
                handler.on_chain_start(
                    {"name": "get_guess"},
                    {"messages": []},
                    run_id=node_a,
                    parent_run_id=root_run,
                )
                node_a_span = hs.runs[str(node_a)]["span"]
                assert node_a_span.span_type == SpanType.CHAIN
                handler.on_chain_end({"messages": []}, run_id=node_a)
                handler.on_chain_end({"messages": []}, run_id=root_run)
                raise Interrupt("Enter a letter:")
        except Interrupt:
            pass

        # Second handoff (resume) — verify node types
        with handler.handoff_context("Hangman") as ctx2:
            assert ctx2._is_resume is True

            # Resume root gets skipped
            resume_root = uuid4()
            handler.on_chain_start(
                {"name": "Hangman Game"},
                {"messages": []},
                run_id=resume_root,
                parent_run_id=orchestrator_tool_run,
            )
            hs = self._get_handoff_state()
            assert str(resume_root) in hs.skipped_resume_roots

            # First node after resume root — reparented, should be CHAIN
            node_b = uuid4()
            handler.on_chain_start(
                {"name": "get_guess"},
                {"messages": []},
                run_id=node_b,
                parent_run_id=resume_root,
            )
            node_b_span = hs.runs[str(node_b)]["span"]
            assert node_b_span.span_type == SpanType.CHAIN

            # End node_b so runs becomes empty
            handler.on_chain_end({"messages": []}, run_id=node_b)
            assert str(node_b) not in hs.runs

            # Second node after resume root — also CHAIN, not AGENT
            node_c = uuid4()
            handler.on_chain_start(
                {"name": "process_guess"},
                {"messages": []},
                run_id=node_c,
                parent_run_id=resume_root,
            )
            node_c_span = hs.runs[str(node_c)]["span"]
            assert node_c_span.span_type == SpanType.CHAIN


class TestGetBeaconHandler:
    """Tests for get_beacon_handler() utility."""

    def test_from_list(self):
        from lumenova_beacon.tracing.integrations.langchain import (
            BeaconLangGraphHandler,
            get_beacon_handler,
        )
        handler = BeaconLangGraphHandler(agent_name="Test")
        config = {"callbacks": [handler]}
        assert get_beacon_handler(config) is handler

    def test_from_callback_manager(self):
        from lumenova_beacon.tracing.integrations.langchain import (
            BeaconLangGraphHandler,
            get_beacon_handler,
        )
        handler = BeaconLangGraphHandler(agent_name="Test")
        manager = MagicMock()
        manager.handlers = [handler]
        config = {"callbacks": manager}
        assert get_beacon_handler(config) is handler

    def test_returns_none_when_absent(self):
        from lumenova_beacon.tracing.integrations.langchain import get_beacon_handler
        config = {"callbacks": [MagicMock()]}
        assert get_beacon_handler(config) is None

    def test_returns_none_when_callbacks_none(self):
        from lumenova_beacon.tracing.integrations.langchain import get_beacon_handler
        config = {"callbacks": None}
        assert get_beacon_handler(config) is None

    def test_returns_none_when_no_callbacks_key(self):
        from lumenova_beacon.tracing.integrations.langchain import get_beacon_handler
        config = {}
        assert get_beacon_handler(config) is None


class TestHandoffSpanType:
    """Tests for the HANDOFF SpanType."""

    def test_handoff_enum_value(self):
        assert SpanType.HANDOFF == "handoff"
        assert SpanType.HANDOFF.value == "handoff"

    def test_create_span_with_handoff_type(self, client):
        span = client.create_span("test", span_type=SpanType.HANDOFF)
        assert span.span_type == SpanType.HANDOFF
        assert span.attributes["span.type"] == "handoff"


class TestAutoDetectSubagentHandoff:
    """Tests for auto-detecting subagent handoffs from nested graph invocations."""

    @pytest.fixture
    def handler(self):
        """Create a BeaconLangGraphHandler with auto-detect enabled."""
        from lumenova_beacon.tracing.integrations.langchain import BeaconLangGraphHandler
        return BeaconLangGraphHandler(
            agent_name="Orchestrator",
            autodetect_handoffs=True,
        )

    @pytest.fixture
    def handler_disabled(self):
        """Create a BeaconLangGraphHandler with auto-detect disabled (default)."""
        from lumenova_beacon.tracing.integrations.langchain import BeaconLangGraphHandler
        return BeaconLangGraphHandler(agent_name="Orchestrator")

    def _get_handoff_state(self):
        """Get the active handoff state from ContextVar."""
        from lumenova_beacon.tracing.integrations.langchain import _active_handoff_state
        return _active_handoff_state.get()

    def _get_effective_trace_id(self, handler):
        """Get the effective trace_id (from ContextVar handoff state or handler base)."""
        hs = self._get_handoff_state()
        return hs.trace_id if hs else handler._current_trace_id

    def _get_effective_runs(self, handler):
        """Get the effective runs dict (from ContextVar handoff state or handler base)."""
        hs = self._get_handoff_state()
        return hs.runs if hs else handler._runs

    def _get_effective_agent_name(self, handler):
        """Get the effective agent_name (from ContextVar handoff state or handler base)."""
        hs = self._get_handoff_state()
        return hs.agent_name if hs else handler._agent_name

    def _setup_active_tool(self, handler, trace_id="parent_trace_aaa"):
        """Set up handler with an active tool span."""
        handler._current_trace_id = trace_id
        handler._root_span_id = "root_span_bbb"
        tool_run_id = uuid4()
        tool_span = Span(name="task", trace_id=trace_id)
        tool_span.start()
        handler._runs = {str(tool_run_id): {"span": tool_span, "parent_run_id": None}}
        handler._active_tool_run_ids = [str(tool_run_id)]
        return tool_run_id

    def _start_subgraph(self, handler, tool_run_id, name="research-assistant", run_id=None):
        """Fire on_chain_start for a subgraph root."""
        run_id = run_id or uuid4()
        handler.on_chain_start(
            serialized={"name": name},
            inputs={"messages": []},
            run_id=run_id,
            parent_run_id=tool_run_id,
            metadata={"langgraph_node": "agent", "langgraph_checkpoint_ns": name},
            tags=["graph:step:1"],
        )
        return run_id

    def test_creates_separate_trace(self, handler):
        """Subgraph root with parent_run_id pointing to active tool should create a new trace."""
        tool_run_id = self._setup_active_tool(handler)
        sub_run_id = self._start_subgraph(handler, tool_run_id)

        # Effective trace (via ContextVar) should be different from parent
        assert self._get_effective_trace_id(handler) != "parent_trace_aaa"
        assert str(sub_run_id) in handler._auto_handoff_contexts

    def test_handoff_span_in_parent_trace(self, handler):
        """HANDOFF span should live in the parent trace."""
        tool_run_id = self._setup_active_tool(handler)
        sub_run_id = self._start_subgraph(handler, tool_run_id)

        ctx = handler._auto_handoff_contexts[str(sub_run_id)]
        assert ctx.handoff_span is not None
        assert ctx.handoff_span.trace_id == "parent_trace_aaa"
        assert ctx.handoff_span.span_type == SpanType.HANDOFF

    def test_handoff_span_attributes_and_link(self, handler):
        """HANDOFF span should have target agent info and OTEL link."""
        tool_run_id = self._setup_active_tool(handler)
        sub_run_id = self._start_subgraph(handler, tool_run_id, name="research-agent")

        ctx = handler._auto_handoff_contexts[str(sub_run_id)]
        assert ctx.handoff_span.attributes["handoff.target_agent"] == "research-agent"
        assert ctx.handoff_span.attributes["handoff.target_trace_id"] == ctx.trace_id
        assert len(ctx.handoff_span.links) == 1
        assert ctx.handoff_span.links[0].attributes == {"link.type": "child_trace"}

    def test_subgraph_root_gets_agent_type(self, handler):
        """First span after auto-handoff should be SpanType.AGENT."""
        tool_run_id = self._setup_active_tool(handler)
        sub_run_id = self._start_subgraph(handler, tool_run_id)

        runs = self._get_effective_runs(handler)
        sub_span = runs[str(sub_run_id)]["span"]
        assert sub_span.span_type == SpanType.AGENT

    def test_restores_state_on_chain_end(self, handler):
        """on_chain_end for subgraph root should restore parent state."""
        tool_run_id = self._setup_active_tool(handler)
        sub_run_id = self._start_subgraph(handler, tool_run_id)

        child_trace = self._get_effective_trace_id(handler)
        assert child_trace != "parent_trace_aaa"

        handler.on_chain_end(outputs={"result": "done"}, run_id=sub_run_id)

        # After exit, ContextVar is cleared, effective state is handler base
        assert self._get_effective_trace_id(handler) == "parent_trace_aaa"
        assert handler._root_span_id == "root_span_bbb"
        assert str(sub_run_id) not in handler._auto_handoff_contexts
        # Tool span should still be in handler base runs
        assert str(tool_run_id) in handler._runs

    def test_restores_state_on_chain_error(self, handler):
        """on_chain_error for subgraph root should restore parent state."""
        tool_run_id = self._setup_active_tool(handler)
        sub_run_id = self._start_subgraph(handler, tool_run_id)

        handler.on_chain_error(
            error=RuntimeError("subgraph crashed"),
            run_id=sub_run_id,
        )

        assert self._get_effective_trace_id(handler) == "parent_trace_aaa"
        assert str(sub_run_id) not in handler._auto_handoff_contexts

    def test_disabled_by_default(self, handler_disabled):
        """Without autodetect_handoffs=True, no handoff is created but span type is AGENT."""
        handler = handler_disabled
        tool_run_id = self._setup_active_tool(handler)
        sub_run_id = self._start_subgraph(handler, tool_run_id)

        # Should stay in same trace (no handoff, no ContextVar set)
        assert self._get_effective_trace_id(handler) == "parent_trace_aaa"
        assert not handler._auto_handoff_contexts
        # Subagent root should still get AGENT type even without full handoff
        runs = self._get_effective_runs(handler)
        span = runs[str(sub_run_id)]["span"]
        assert span.span_type == SpanType.AGENT

    def test_no_handoff_without_langgraph_metadata(self, handler):
        """Without LangGraph metadata, no handoff even with parent pointing to tool."""
        tool_run_id = self._setup_active_tool(handler)
        sub_run_id = uuid4()

        handler.on_chain_start(
            serialized={"name": "plain-chain"},
            inputs={},
            run_id=sub_run_id,
            parent_run_id=tool_run_id,
            metadata={},
            tags=[],
        )

        assert self._get_effective_trace_id(handler) == "parent_trace_aaa"
        assert not handler._auto_handoff_contexts
        runs = self._get_effective_runs(handler)
        span = runs[str(sub_run_id)]["span"]
        assert span.span_type == SpanType.CHAIN

    def test_no_handoff_when_parent_is_not_tool(self, handler):
        """Chain whose parent is NOT an active tool should not trigger handoff."""
        self._setup_active_tool(handler)
        non_tool_parent = uuid4()
        sub_run_id = uuid4()

        handler.on_chain_start(
            serialized={"name": "inner-chain"},
            inputs={},
            run_id=sub_run_id,
            parent_run_id=non_tool_parent,
            metadata={"langgraph_node": "agent"},
            tags=["graph:step:1"],
        )

        assert self._get_effective_trace_id(handler) == "parent_trace_aaa"
        assert not handler._auto_handoff_contexts

    def test_handoff_span_exported_on_exit(self, handler):
        """HANDOFF span should be exported when subgraph root ends."""
        tool_run_id = self._setup_active_tool(handler)
        sub_run_id = self._start_subgraph(handler, tool_run_id)

        with patch.object(handler.client, "export_span") as mock_export:
            handler.on_chain_end(outputs={}, run_id=sub_run_id)

        exported_types = [call.args[0].span_type for call in mock_export.call_args_list]
        assert SpanType.HANDOFF in exported_types

    def test_nested_auto_handoffs(self, handler):
        """Subagent A invokes Subagent B — both should auto-handoff correctly."""
        tool_a_run = self._setup_active_tool(handler)
        original_trace = handler._current_trace_id

        # Subagent A enters (parent_run_id = tool_a)
        sub_a_run = self._start_subgraph(handler, tool_a_run, name="agent-a")
        trace_a = self._get_effective_trace_id(handler)
        assert trace_a != original_trace

        # Tool inside agent-a that will invoke subagent B
        tool_b_run = uuid4()
        handler.on_tool_start(
            serialized={"name": "task"},
            input_str="",
            run_id=tool_b_run,
            parent_run_id=sub_a_run,
        )

        # Subagent B enters (parent_run_id = tool_b)
        sub_b_run = self._start_subgraph(handler, tool_b_run, name="agent-b")
        trace_b = self._get_effective_trace_id(handler)
        assert trace_b != trace_a
        assert trace_b != original_trace

        # End agent-b
        handler.on_chain_end(outputs={}, run_id=sub_b_run)
        assert self._get_effective_trace_id(handler) == trace_a

        # End tool_b
        handler.on_tool_end(output="done", run_id=tool_b_run)

        # End agent-a
        handler.on_chain_end(outputs={}, run_id=sub_a_run)
        assert self._get_effective_trace_id(handler) == original_trace

    def test_agent_name_set_during_handoff(self, handler):
        """agent_name should be the subagent name during auto-handoff."""
        tool_run_id = self._setup_active_tool(handler)
        assert handler._agent_name == "Orchestrator"

        sub_run_id = self._start_subgraph(handler, tool_run_id, name="research-assistant")
        assert self._get_effective_agent_name(handler) == "research-assistant"

        handler.on_chain_end(outputs={}, run_id=sub_run_id)
        assert self._get_effective_agent_name(handler) == "Orchestrator"

    def test_parallel_auto_handoffs_parent_under_correct_tools(self, handler):
        """Parallel auto-handoffs should parent under their respective tool spans."""
        handler._current_trace_id = "parent_trace"
        handler._root_span_id = "root_span"

        # Set up 3 active tool spans (simulating parallel tool execution)
        tool_run_ids = [uuid4() for _ in range(3)]
        tool_spans = []
        for i, trid in enumerate(tool_run_ids):
            tool_span = Span(name=f"handoff_to_agent_{i}", trace_id="parent_trace")
            tool_span.start()
            tool_spans.append(tool_span)
            handler._runs[str(trid)] = {"span": tool_span, "parent_run_id": None}
        handler._active_tool_run_ids = [str(trid) for trid in tool_run_ids]

        # Start subgraphs sequentially, exiting each before the next
        # (simulates what happens per-task in async parallel execution)
        handoff_parents: list[str] = []
        for i, trid in enumerate(tool_run_ids):
            sub_run_id = self._start_subgraph(handler, trid, name=f"agent-{i}")
            ctx = handler._auto_handoff_contexts[str(sub_run_id)]
            handoff_parents.append(ctx.handoff_span.parent_id)
            # Exit handoff so next iteration sees parent state
            handler.on_chain_end(outputs={}, run_id=sub_run_id)

        # Each handoff span should parent under its respective tool span
        for i, parent_id in enumerate(handoff_parents):
            assert parent_id == tool_spans[i].span_id, (
                f"Handoff {i} should parent under tool {i} "
                f"(expected {tool_spans[i].span_id}, got {parent_id})"
            )

    def test_invocation_completed_not_set_by_auto_handoff_child(self, handler):
        """Auto-handoff child root ending should NOT set _invocation_completed."""
        tool_run_id = self._setup_active_tool(handler)
        assert handler._invocation_completed is False

        sub_run_id = self._start_subgraph(handler, tool_run_id)

        # End the subgraph root (this is a root span with parent_run_id=None)
        handler.on_chain_end(outputs={"result": "done"}, run_id=sub_run_id)

        # _invocation_completed should still be False because the child
        # trace's root ending should not mark the handler's invocation complete
        assert handler._invocation_completed is False
