"""Tests for resume output re-export on root span."""

import json
from unittest.mock import patch, MagicMock
from uuid import uuid4

import pytest

from lumenova_beacon.tracing.integrations.langchain import BeaconLangGraphHandler
from lumenova_beacon.types import SpanType, StatusCode


class TestUpdateRootSpanWithResumeOutput:
    """Tests for _update_root_span_with_resume_output()."""

    @pytest.fixture
    def handler(self):
        """Create a BeaconLangGraphHandler configured as a resume."""
        h = BeaconLangGraphHandler(agent_name="Conversation Agent")
        h._current_trace_id = "a" * 32
        h._root_span_id = "b" * 16
        h._session_id = "test-session"
        return h

    def test_exports_span_with_same_ids(self, handler):
        """Re-exported span should use the same trace_id and span_id as root."""
        with patch.object(handler.client, "export_span") as mock_export:
            handler._update_root_span_with_resume_output({"messages": ["hello"]})

            mock_export.assert_called_once()
            span = mock_export.call_args[0][0]
            assert span.trace_id == "a" * 32
            assert span.span_id == "b" * 16

    def test_span_type_is_agent(self, handler):
        """Re-exported span should have AGENT type."""
        with patch.object(handler.client, "export_span") as mock_export:
            handler._update_root_span_with_resume_output({"messages": ["hello"]})

            span = mock_export.call_args[0][0]
            assert span.span_type == SpanType.AGENT

    def test_span_has_output(self, handler):
        """Re-exported span should contain the final output."""
        outputs = {"messages": [{"role": "assistant", "content": "Answer is 42"}]}
        with patch.object(handler.client, "export_span") as mock_export:
            handler._update_root_span_with_resume_output(outputs)

            span = mock_export.call_args[0][0]
            output_attr = span.attributes.get("gen_ai.output.messages")
            assert output_attr is not None
            parsed = json.loads(output_attr)
            assert parsed["messages"][0]["content"] == "Answer is 42"

    def test_span_has_agent_name(self, handler):
        """Re-exported span should carry the agent name."""
        with patch.object(handler.client, "export_span") as mock_export:
            handler._update_root_span_with_resume_output({"messages": []})

            span = mock_export.call_args[0][0]
            assert span.attributes["gen_ai.agent.name"] == "Conversation Agent"

    def test_span_overrides_cancelled(self, handler):
        """Re-exported span should set langgraph.cancelled=False and langgraph.resumed=True."""
        with patch.object(handler.client, "export_span") as mock_export:
            handler._update_root_span_with_resume_output({"messages": []})

            span = mock_export.call_args[0][0]
            assert span.attributes["langgraph.cancelled"] is False
            assert span.attributes["langgraph.resumed"] is True

    def test_span_status_ok(self, handler):
        """Re-exported span should have OK status."""
        with patch.object(handler.client, "export_span") as mock_export:
            handler._update_root_span_with_resume_output({"messages": []})

            span = mock_export.call_args[0][0]
            assert span.status_code == StatusCode.OK

    def test_span_has_session_id(self, handler):
        """Re-exported span should carry the session_id."""
        with patch.object(handler.client, "export_span") as mock_export:
            handler._update_root_span_with_resume_output({"messages": []})

            span = mock_export.call_args[0][0]
            assert span.session_id == "test-session"

    def test_noop_when_no_root_span_id(self, handler):
        """Should not export if _root_span_id is not set."""
        handler._root_span_id = None
        with patch.object(handler.client, "export_span") as mock_export:
            handler._update_root_span_with_resume_output({"messages": []})
            mock_export.assert_not_called()

    def test_noop_when_no_trace_id(self, handler):
        """Should not export if _current_trace_id is not set."""
        handler._current_trace_id = None
        with patch.object(handler.client, "export_span") as mock_export:
            handler._update_root_span_with_resume_output({"messages": []})
            mock_export.assert_not_called()


class TestOnChainEndResumeRoot:
    """Tests for on_chain_end handling of skipped resume roots."""

    @pytest.fixture
    def handler(self):
        """Create a handler simulating a resume state."""
        h = BeaconLangGraphHandler(agent_name="Conversation Agent")
        h._current_trace_id = "a" * 32
        h._root_span_id = "b" * 16
        h._is_resume = True
        return h

    def test_resume_root_triggers_re_export(self, handler):
        """on_chain_end for a skipped resume root should call _update_root_span_with_resume_output."""
        run_id = uuid4()
        handler._skipped_resume_roots[str(run_id)] = {
            "parent_id": None,
            "resume_value": None,
        }

        with patch.object(handler, "_update_root_span_with_resume_output") as mock_update:
            handler.on_chain_end(
                outputs={"messages": ["final answer"]},
                run_id=run_id,
            )
            mock_update.assert_called_once_with({"messages": ["final answer"]})

    def test_resume_root_cleaned_up_after(self, handler):
        """Skipped resume root entry should be removed after on_chain_end."""
        run_id = uuid4()
        handler._skipped_resume_roots[str(run_id)] = {
            "parent_id": None,
            "resume_value": None,
        }

        with patch.object(handler, "_update_root_span_with_resume_output"):
            handler.on_chain_end(
                outputs={"messages": []},
                run_id=run_id,
            )
        assert str(run_id) not in handler._skipped_resume_roots

    def test_resume_root_does_not_call_end_span(self, handler):
        """on_chain_end for resume root should NOT call _end_span (no run in _runs)."""
        run_id = uuid4()
        handler._skipped_resume_roots[str(run_id)] = {
            "parent_id": None,
            "resume_value": None,
        }

        with patch.object(handler, "_update_root_span_with_resume_output"):
            with patch.object(handler, "_end_span") as mock_end:
                handler.on_chain_end(
                    outputs={"messages": []},
                    run_id=run_id,
                )
                mock_end.assert_not_called()
