"""Tests for trace context management."""

from unittest.mock import MagicMock, patch

import pytest

from lumenova_beacon.tracing.span import Span
from lumenova_beacon.tracing.trace import (
    TraceContext,
    get_current_trace_id,
    get_current_span,
    set_current_trace_id,
    set_current_span,
    clear_context,
)
from lumenova_beacon.types import SpanType, StatusCode


class BrokenStrException(Exception):
    """Exception whose __str__ raises, simulating complex framework exceptions."""
    def __str__(self):
        raise RuntimeError("broken __str__")


@pytest.fixture
def mock_client():
    return MagicMock()


class TestContextFunctions:
    """Test context variable functions."""

    def test_get_current_trace_id_when_none(self):
        """Test getting trace ID when none is set."""
        clear_context()

        trace_id = get_current_trace_id()

        assert trace_id is None

    def test_set_and_get_current_trace_id(self):
        """Test setting and getting trace ID."""
        set_current_trace_id("trace-123")

        trace_id = get_current_trace_id()

        assert trace_id == "trace-123"

    def test_get_current_span_when_none(self):
        """Test getting span when none is set."""
        clear_context()

        span = get_current_span()

        assert span is None

    def test_set_and_get_current_span(self):
        """Test setting and getting current span."""
        test_span = Span(name="test", trace_id="trace-123")

        set_current_span(test_span)

        span = get_current_span()

        assert span is test_span

    def test_set_current_span_also_sets_trace_id(self):
        """Test that setting span also sets trace ID."""
        test_span = Span(name="test", trace_id="trace-123")

        set_current_span(test_span)

        trace_id = get_current_trace_id()

        assert trace_id == "trace-123"

    def test_clear_context(self):
        """Test clearing the context."""
        test_span = Span(name="test", trace_id="trace-123")
        set_current_span(test_span)

        clear_context()

        assert get_current_span() is None
        assert get_current_trace_id() is None


class TestTraceContextInitialization:
    """Test TraceContext initialization."""

    def test_init_with_span_and_client(self, mock_client):
        """Test initialization with span and client."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        assert ctx.span is span
        assert ctx.client is mock_client
        assert ctx.previous_span is None
        assert ctx.previous_trace_id is None


class TestTraceContextSyncContextManager:
    """Test TraceContext as synchronous context manager."""

    def test_enter_sets_current_span(self, mock_client):
        """Test that entering context sets the current span."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with ctx:
            current = get_current_span()
            assert current is span

    def test_enter_starts_span(self, mock_client):
        """Test that entering context starts the span."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with ctx:
            assert span.start_time is not None

    def test_enter_doesnt_restart_started_span(self, mock_client):
        """Test that already-started span is not restarted."""
        span = Span(name="test")
        span.start()
        first_start_time = span.start_time

        ctx = TraceContext(span, client=mock_client)

        with ctx:
            # Start time should remain the same
            assert span.start_time == first_start_time

    def test_enter_returns_span(self, mock_client):
        """Test that entering context returns the span."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with ctx as s:
            assert s is span

    def test_exit_ends_span(self, mock_client):
        """Test that exiting context ends the span."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with ctx:
            pass

        assert span.end_time is not None

    def test_exit_sets_ok_status_on_success(self, mock_client):
        """Test that successful exit sets OK status."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with ctx:
            pass

        assert span.status_code == StatusCode.OK

    def test_exit_doesnt_override_existing_status(self, mock_client):
        """Test that exit doesn't override already-set status."""
        span = Span(name="test")
        span.set_status(StatusCode.ERROR)
        ctx = TraceContext(span, client=mock_client)

        with ctx:
            pass

        # Should keep ERROR status
        assert span.status_code == StatusCode.ERROR

    def test_exit_records_exception(self, mock_client):
        """Test that exit records exceptions."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        try:
            with ctx:
                raise ValueError("Test error")
        except ValueError:
            pass

        assert span.status_code == StatusCode.ERROR
        assert span.status_description == "Test error"
        assert span.attributes["exception.type"] == "ValueError"

    def test_exit_propagates_exception(self, mock_client):
        """Test that exceptions are propagated."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with pytest.raises(ValueError):
            with ctx:
                raise ValueError("Test error")

    def test_exit_restores_previous_context(self, mock_client):
        """Test that exit restores previous context."""
        parent_span = Span(name="parent", trace_id="trace-parent")
        set_current_span(parent_span)

        child_span = Span(name="child", trace_id="trace-child")
        ctx = TraceContext(child_span, client=mock_client)

        with ctx:
            # Inside context, should have child span
            assert get_current_span() is child_span

        # After context, should restore parent span
        assert get_current_span() is parent_span

    def test_exit_calls_client_export_span(self, mock_client):
        """Test that exit exports span via client."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with ctx:
            pass

        mock_client.export_span.assert_called_once_with(span)


class TestTraceContextAsyncContextManager:
    """Test TraceContext as asynchronous context manager."""

    @pytest.mark.asyncio
    async def test_aenter_sets_current_span(self, mock_client):
        """Test that entering async context sets the current span."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        async with ctx:
            current = get_current_span()
            assert current is span

    @pytest.mark.asyncio
    async def test_aenter_starts_span(self, mock_client):
        """Test that entering async context starts the span."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        async with ctx:
            assert span.start_time is not None

    @pytest.mark.asyncio
    async def test_aenter_returns_span(self, mock_client):
        """Test that entering async context returns the span."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        async with ctx as s:
            assert s is span

    @pytest.mark.asyncio
    async def test_aexit_ends_span(self, mock_client):
        """Test that exiting async context ends the span."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        async with ctx:
            pass

        assert span.end_time is not None

    @pytest.mark.asyncio
    async def test_aexit_sets_ok_status_on_success(self, mock_client):
        """Test that successful async exit sets OK status."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        async with ctx:
            pass

        assert span.status_code == StatusCode.OK

    @pytest.mark.asyncio
    async def test_aexit_records_exception(self, mock_client):
        """Test that async exit records exceptions."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        try:
            async with ctx:
                raise ValueError("Test error")
        except ValueError:
            pass

        assert span.status_code == StatusCode.ERROR
        assert span.status_description == "Test error"

    @pytest.mark.asyncio
    async def test_aexit_propagates_exception(self, mock_client):
        """Test that exceptions are propagated in async context."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with pytest.raises(ValueError):
            async with ctx:
                raise ValueError("Test error")

    @pytest.mark.asyncio
    async def test_aexit_restores_previous_context(self, mock_client):
        """Test that async exit restores previous context."""
        parent_span = Span(name="parent", trace_id="trace-parent")
        set_current_span(parent_span)

        child_span = Span(name="child", trace_id="trace-child")
        ctx = TraceContext(child_span, client=mock_client)

        async with ctx:
            # Inside context, should have child span
            assert get_current_span() is child_span

        # After context, should restore parent span
        assert get_current_span() is parent_span

    @pytest.mark.asyncio
    async def test_aexit_calls_client_export_span(self, mock_client):
        """Test that async exit exports span via client."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        async with ctx:
            pass

        mock_client.export_span.assert_called_once_with(span)


class TestTraceContextNesting:
    """Test nested TraceContext scenarios."""

    def test_nested_contexts(self, mock_client):
        """Test nested trace contexts."""
        parent_span = Span(name="parent")
        child_span = Span(name="child")

        parent_ctx = TraceContext(parent_span, client=mock_client)
        child_ctx = TraceContext(child_span, client=mock_client)

        with parent_ctx:
            assert get_current_span() is parent_span

            with child_ctx:
                assert get_current_span() is child_span

            # Back to parent
            assert get_current_span() is parent_span

        # Context cleared
        assert get_current_span() is None

    def test_nested_contexts_with_trace_id_only(self, mock_client):
        """Test nested contexts with trace ID but no parent span."""
        set_current_trace_id("trace-123")

        child_span = Span(name="child")
        ctx = TraceContext(child_span, client=mock_client)

        with ctx:
            assert get_current_span() is child_span

        # Should restore trace ID only
        assert get_current_trace_id() == "trace-123"
        assert get_current_span() is None

    def test_deeply_nested_contexts(self, mock_client):
        """Test deeply nested trace contexts."""
        span1 = Span(name="level1")
        span2 = Span(name="level2")
        span3 = Span(name="level3")

        with TraceContext(span1, client=mock_client):
            assert get_current_span() is span1

            with TraceContext(span2, client=mock_client):
                assert get_current_span() is span2

                with TraceContext(span3, client=mock_client):
                    assert get_current_span() is span3

                assert get_current_span() is span2

            assert get_current_span() is span1


class TestTraceContextEdgeCases:
    """Test TraceContext edge cases."""

    def test_exit_doesnt_end_already_ended_span(self, mock_client):
        """Test that exit doesn't re-end an already-ended span."""
        span = Span(name="test")
        span.start()
        span.end()
        first_end_time = span.end_time

        ctx = TraceContext(span, client=mock_client)

        with ctx:
            pass

        # End time should remain the same
        assert span.end_time == first_end_time

    def test_context_with_exception_still_restores(self, mock_client):
        """Test that context is restored even when exception occurs."""
        parent_span = Span(name="parent")
        set_current_span(parent_span)

        child_span = Span(name="child")
        ctx = TraceContext(child_span, client=mock_client)

        try:
            with ctx:
                raise ValueError("Test")
        except ValueError:
            pass

        # Should still restore parent context
        assert get_current_span() is parent_span

    def test_restore_context_with_no_previous_context(self, mock_client):
        """Test restoring context when there was no previous context."""
        clear_context()

        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with ctx:
            assert get_current_span() is span

        # Should clear context
        assert get_current_span() is None
        assert get_current_trace_id() is None

    @pytest.mark.asyncio
    async def test_async_context_with_no_previous_context(self, mock_client):
        """Test async context when there was no previous context."""
        clear_context()

        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        async with ctx:
            assert get_current_span() is span

        # Should clear context
        assert get_current_span() is None
        assert get_current_trace_id() is None

    def test_multiple_sequential_contexts(self, mock_client):
        """Test multiple sequential (non-nested) contexts."""
        span1 = Span(name="test1")
        span2 = Span(name="test2")
        span3 = Span(name="test3")

        with TraceContext(span1, client=mock_client):
            assert get_current_span() is span1

        assert get_current_span() is None

        with TraceContext(span2, client=mock_client):
            assert get_current_span() is span2

        assert get_current_span() is None

        with TraceContext(span3, client=mock_client):
            assert get_current_span() is span3

        assert get_current_span() is None


class TestTraceContextErrorResilience:
    """Test that TraceContext always finishes spans even when cleanup code raises."""

    def test_span_gets_end_time_when_record_exception_raises(self, mock_client):
        """Span must have end_time even when record_exception raises."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        try:
            with ctx:
                raise BrokenStrException("boom")
        except (BrokenStrException, RuntimeError):
            pass

        assert span.end_time is not None

    def test_context_restored_when_record_exception_raises(self, mock_client):
        """Previous context must be restored when record_exception raises."""
        clear_context()
        parent_span = Span(name="parent", trace_id="trace-parent")
        set_current_span(parent_span)

        child_span = Span(name="child", trace_id="trace-child")
        ctx = TraceContext(child_span, client=mock_client)

        try:
            with ctx:
                raise BrokenStrException("boom")
        except (BrokenStrException, RuntimeError):
            pass

        assert get_current_span() is parent_span

    def test_context_restored_when_span_end_raises(self, mock_client):
        """Context must be restored even when span.end() raises."""
        clear_context()
        parent_span = Span(name="parent", trace_id="trace-parent")
        set_current_span(parent_span)

        child_span = Span(name="child", trace_id="trace-child")
        ctx = TraceContext(child_span, client=mock_client)

        with patch.object(child_span, "end", side_effect=RuntimeError("end failed")):
            with ctx:
                pass

        assert get_current_span() is parent_span

    def test_context_restored_when_export_span_raises(self, mock_client):
        """Context must be restored even when _export_span raises."""
        clear_context()
        parent_span = Span(name="parent", trace_id="trace-parent")
        set_current_span(parent_span)

        child_span = Span(name="child", trace_id="trace-child")
        ctx = TraceContext(child_span, client=mock_client)

        with patch.object(ctx, "_export_span", side_effect=RuntimeError("export failed")):
            with ctx:
                pass

        assert get_current_span() is parent_span

    def test_span_still_exported_when_record_exception_raises(self, mock_client):
        """_export_span must still be called even when record_exception raises."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with patch.object(ctx, "_export_span") as mock_export:
            try:
                with ctx:
                    raise BrokenStrException("boom")
            except (BrokenStrException, RuntimeError):
                pass

        mock_export.assert_called_once()

    def test_original_exception_propagates_despite_cleanup_errors(self, mock_client):
        """The user's original exception must propagate even if cleanup fails."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with pytest.raises(ValueError, match="user error"):
            with ctx:
                raise ValueError("user error")

    def test_span_gets_error_status_when_record_exception_raises(self, mock_client):
        """Span status should be ERROR with exception details when __str__ raises."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        try:
            with ctx:
                raise BrokenStrException("boom")
        except (BrokenStrException, RuntimeError):
            pass

        assert span.status_code == StatusCode.ERROR
        assert span.attributes["exception.type"] == "BrokenStrException"
        assert span.attributes.get("exception.message") is not None

    @pytest.mark.asyncio
    async def test_async_span_gets_end_time_when_record_exception_raises(self, mock_client):
        """Async: span must have end_time even when record_exception raises."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        try:
            async with ctx:
                raise BrokenStrException("boom")
        except (BrokenStrException, RuntimeError):
            pass

        assert span.end_time is not None

    @pytest.mark.asyncio
    async def test_async_context_restored_when_record_exception_raises(self, mock_client):
        """Async: previous context must be restored when record_exception raises."""
        clear_context()
        parent_span = Span(name="parent", trace_id="trace-parent")
        set_current_span(parent_span)

        child_span = Span(name="child", trace_id="trace-child")
        ctx = TraceContext(child_span, client=mock_client)

        try:
            async with ctx:
                raise BrokenStrException("boom")
        except (BrokenStrException, RuntimeError):
            pass

        assert get_current_span() is parent_span

    @pytest.mark.asyncio
    async def test_async_span_gets_error_status_when_record_exception_raises(self, mock_client):
        """Async: span status should be ERROR when record_exception raises."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        try:
            async with ctx:
                raise BrokenStrException("boom")
        except (BrokenStrException, RuntimeError):
            pass

        assert span.status_code == StatusCode.ERROR


class TestEagerExport:
    """Test eager span export behavior."""

    def test_eager_export_called_on_enter(self, mock_client):
        """Eager export is delegated to client on __enter__."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with ctx:
            pass

        mock_client.export_eager_span.assert_called_once_with(span)

    def test_eager_export_failure_does_not_break_span_lifecycle(self, mock_client):
        """Eager export failure must never break span execution."""
        mock_client.export_eager_span.side_effect = RuntimeError("export failed")
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        with ctx:
            pass

        # Span should still complete normally
        assert span.end_time is not None
        assert span.status_code == StatusCode.OK

    def test_eager_exported_span_has_end_time_zero(self):
        """Eager-exported span has end_time=0 (OTEL incomplete convention)."""
        span = Span(name="test")
        span.start()

        otel_span = span.to_otel_span_eager()

        assert otel_span.end_time == 0

    def test_eager_exported_span_has_unset_status(self):
        """Eager-exported span has status=UNSET."""
        from opentelemetry.trace.status import StatusCode as OTelStatusCode

        span = Span(name="test")
        span.start()

        otel_span = span.to_otel_span_eager()

        assert otel_span.status.status_code == OTelStatusCode.UNSET

    def test_eager_exported_span_includes_attributes(self):
        """Eager-exported span includes currently-set attributes."""
        span = Span(name="test", span_type=SpanType.GENERATION, session_id="sess-123")
        span.set_attribute("code.function", "my_func")
        span.start()

        otel_span = span.to_otel_span_eager()

        assert otel_span.attributes["span.type"] == "generation"
        assert otel_span.attributes["session_id"] == "sess-123"
        assert otel_span.attributes["code.function"] == "my_func"

    def test_eager_and_final_export_use_same_span_id(self):
        """Both eager and final export use the same span_id."""
        span = Span(name="test")
        span.start()

        eager_otel = span.to_otel_span_eager()
        span.end()
        final_otel = span.to_otel_span()

        assert eager_otel.context.span_id == final_otel.context.span_id
        assert eager_otel.context.trace_id == final_otel.context.trace_id

    @pytest.mark.asyncio
    async def test_eager_export_called_on_aenter(self, mock_client):
        """Eager export is delegated to client on __aenter__."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        async with ctx:
            pass

        mock_client.export_eager_span.assert_called_once_with(span)

    @pytest.mark.asyncio
    async def test_eager_export_failure_does_not_break_async_span(self, mock_client):
        """Eager export failure must never break async span execution."""
        mock_client.export_eager_span.side_effect = RuntimeError("export failed")
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        async with ctx:
            pass

        assert span.end_time is not None
        assert span.status_code == StatusCode.OK

    def test_eager_export_with_exception_still_works(self, mock_client):
        """Eager export is called even when the body raises."""
        span = Span(name="test")
        ctx = TraceContext(span, client=mock_client)

        try:
            with ctx:
                raise ValueError("user error")
        except ValueError:
            pass

        # Eager export was called during __enter__
        mock_client.export_eager_span.assert_called_once_with(span)
        # Span still records the error
        assert span.status_code == StatusCode.ERROR
