"""Tests for Flowlines context attributes (context manager + imperative API)."""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING
from unittest.mock import MagicMock, patch

import pytest
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
from opentelemetry.sdk.trace.export import (
    SimpleSpanProcessor,
    SpanExporter,
    SpanExportResult,
)

if TYPE_CHECKING:
    from collections.abc import Sequence

from flowlines._context import FlowlinesSpanProcessor, _flowlines_attrs
from flowlines._init import _reset, clear_context, context, init, set_context


class _CollectingExporter(SpanExporter):
    """Test exporter that collects spans in a list."""

    def __init__(self) -> None:
        self.spans: list[ReadableSpan] = []

    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
        self.spans.extend(spans)
        return SpanExportResult.SUCCESS

    def shutdown(self) -> None:
        pass

    def force_flush(self, timeout_millis: int = 30000) -> bool:
        return True


@pytest.fixture(autouse=True)
def _reset_singleton() -> None:
    """Reset the Flowlines singleton before each test."""
    _reset()


def _make_provider_and_exporter() -> tuple[TracerProvider, _CollectingExporter]:
    """Create a TracerProvider with FlowlinesSpanProcessor for testing."""
    exporter = _CollectingExporter()
    inner_processor = SimpleSpanProcessor(exporter)
    processor = FlowlinesSpanProcessor(inner_processor)
    provider = TracerProvider()
    provider.add_span_processor(processor)
    return provider, exporter


class TestContextManager:
    @patch("flowlines._init.atexit", autospec=True)
    def test_injects_attributes_into_spans(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """Context manager injects user_id and session_id into spans."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        with (
            context(user_id="user-42", session_id="convo-abc"),
            tracer.start_as_current_span("test-span"),
        ):
            pass

        provider.shutdown()
        spans = exporter.spans
        assert len(spans) == 1
        attrs = dict(spans[0].attributes or {})
        assert attrs["flowlines.user_id"] == "user-42"
        assert attrs["flowlines.session_id"] == "convo-abc"

    @patch("flowlines._init.atexit", autospec=True)
    def test_only_mandatory_attrs(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """Context manager with only mandatory attrs sets only those attributes."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        with (
            context(user_id="user-42", session_id="sess-1"),
            tracer.start_as_current_span("test-span"),
        ):
            pass

        provider.shutdown()
        spans = exporter.spans
        attrs = dict(spans[0].attributes or {})
        assert attrs["flowlines.user_id"] == "user-42"
        assert attrs["flowlines.session_id"] == "sess-1"
        assert "flowlines.agent_id" not in attrs

    @patch("flowlines._init.atexit", autospec=True)
    def test_no_context_no_attributes(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """Without context, spans have no flowlines.* attributes."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        with tracer.start_as_current_span("test-span"):
            pass

        provider.shutdown()
        spans = exporter.spans
        attrs = dict(spans[0].attributes or {})
        assert "flowlines.user_id" not in attrs
        assert "flowlines.session_id" not in attrs


class TestNesting:
    @patch("flowlines._init.atexit", autospec=True)
    def test_inner_replaces_outer(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """Inner context fully replaces outer (no merge)."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        with (
            context(user_id="outer-user", session_id="outer-convo"),
            context(user_id="inner-user", session_id="inner-convo"),
            tracer.start_as_current_span("inner-span"),
        ):
            pass

        provider.shutdown()
        spans = exporter.spans
        attrs = dict(spans[0].attributes or {})
        assert attrs["flowlines.user_id"] == "inner-user"
        # session_id from outer is NOT inherited
        assert attrs["flowlines.session_id"] == "inner-convo"

    @patch("flowlines._init.atexit", autospec=True)
    def test_outer_restored_after_inner(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """After inner context exits, outer context is restored."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        with context(user_id="outer-user", session_id="outer-convo"):
            with context(user_id="inner-user", session_id="inner-convo"):
                pass
            # After inner exits, outer should be active again
            with tracer.start_as_current_span("after-inner"):
                pass

        provider.shutdown()
        spans = exporter.spans
        attrs = dict(spans[0].attributes or {})
        assert attrs["flowlines.user_id"] == "outer-user"
        assert attrs["flowlines.session_id"] == "outer-convo"


class TestExceptionSafety:
    @patch("flowlines._init.atexit", autospec=True)
    def test_context_restored_on_exception(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """Context is restored even if an exception occurs inside the block."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )

        with context(user_id="outer", session_id="sess-outer"):
            try:
                with context(user_id="inner", session_id="sess-inner"):
                    raise ValueError("boom")
            except ValueError:
                pass
            # Outer context should be restored
            assert _flowlines_attrs.get() == {
                "flowlines.user_id": "outer",
                "flowlines.session_id": "sess-outer",
            }

        # After all contexts exit, should be None (default)
        assert _flowlines_attrs.get() is None


class TestImperativeAPI:
    @patch("flowlines._init.atexit", autospec=True)
    def test_set_and_clear_context(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """set_context/clear_context injects attributes into spans."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        token = set_context(user_id="user-99", session_id="convo-xyz")
        try:
            with tracer.start_as_current_span("test-span"):
                pass
        finally:
            clear_context(token)

        provider.shutdown()
        spans = exporter.spans
        attrs = dict(spans[0].attributes or {})
        assert attrs["flowlines.user_id"] == "user-99"
        assert attrs["flowlines.session_id"] == "convo-xyz"

    @patch("flowlines._init.atexit", autospec=True)
    def test_clear_restores_previous(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """clear_context restores the previous context state."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )

        token = set_context(user_id="user-1", session_id="sess-1")
        assert _flowlines_attrs.get() == {
            "flowlines.user_id": "user-1",
            "flowlines.session_id": "sess-1",
        }
        clear_context(token)
        assert _flowlines_attrs.get() is None


class TestAgentId:
    @patch("flowlines._init.atexit", autospec=True)
    def test_context_injects_agent_id(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """Context manager injects user_id, session_id, and agent_id into spans."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        with (
            context(user_id="user-42", session_id="sess-abc", agent_id="agent-1"),
            tracer.start_as_current_span("test-span"),
        ):
            pass

        provider.shutdown()
        spans = exporter.spans
        assert len(spans) == 1
        attrs = dict(spans[0].attributes or {})
        assert attrs["flowlines.user_id"] == "user-42"
        assert attrs["flowlines.session_id"] == "sess-abc"
        assert attrs["flowlines.agent_id"] == "agent-1"

    @patch("flowlines._init.atexit", autospec=True)
    def test_agent_id_omitted(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """When agent_id is not passed, the attribute is not present on spans."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        with (
            context(user_id="user-42", session_id="sess-1"),
            tracer.start_as_current_span("test-span"),
        ):
            pass

        provider.shutdown()
        spans = exporter.spans
        attrs = dict(spans[0].attributes or {})
        assert attrs["flowlines.user_id"] == "user-42"
        assert attrs["flowlines.session_id"] == "sess-1"
        assert "flowlines.agent_id" not in attrs

    @patch("flowlines._init.atexit", autospec=True)
    def test_imperative_set_context_with_agent_id(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """set_context with agent_id injects the attribute into spans."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        token = set_context(
            user_id="user-99", session_id="sess-xyz", agent_id="agent-2"
        )
        try:
            with tracer.start_as_current_span("test-span"):
                pass
        finally:
            clear_context(token)

        provider.shutdown()
        spans = exporter.spans
        attrs = dict(spans[0].attributes or {})
        assert attrs["flowlines.user_id"] == "user-99"
        assert attrs["flowlines.session_id"] == "sess-xyz"
        assert attrs["flowlines.agent_id"] == "agent-2"


class TestAsyncSafety:
    @patch("flowlines._init.atexit", autospec=True)
    def test_contextvars_isolation_between_tasks(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """Context attributes are isolated between concurrent async tasks."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        async def task_a() -> None:
            with context(user_id="user-a", session_id="sess-a"):
                await asyncio.sleep(0.01)
                with tracer.start_as_current_span("span-a"):
                    pass

        async def task_b() -> None:
            with context(user_id="user-b", session_id="sess-b"):
                await asyncio.sleep(0.01)
                with tracer.start_as_current_span("span-b"):
                    pass

        async def run() -> None:
            await asyncio.gather(task_a(), task_b())

        asyncio.run(run())
        provider.shutdown()

        spans = exporter.spans
        assert len(spans) == 2

        span_map = {s.name: dict(s.attributes or {}) for s in spans}
        assert span_map["span-a"]["flowlines.user_id"] == "user-a"
        assert span_map["span-b"]["flowlines.user_id"] == "user-b"


class TestNestedSpanPropagation:
    @patch("flowlines._init.atexit", autospec=True)
    def test_context_attributes_on_all_nested_spans(
        self,
        _mock_atexit: MagicMock,
    ) -> None:
        """Flowlines attributes are present on every nested span."""
        init(
            api_key="key",
            ingest_endpoint="https://example.com",
            has_external_otel=True,
        )
        provider, exporter = _make_provider_and_exporter()
        tracer = provider.get_tracer("test")

        with (
            context(user_id="user-42", session_id="sess-abc"),
            tracer.start_as_current_span("root-span"),
            tracer.start_as_current_span("middle-span"),
            tracer.start_as_current_span("leaf-span") as leaf,
        ):
            leaf.set_attribute("gen_ai.system", "openai")

        provider.shutdown()
        spans = exporter.spans
        assert len(spans) == 3

        for span in spans:
            attrs = dict(span.attributes or {})
            assert attrs["flowlines.user_id"] == "user-42", (
                f"span '{span.name}' missing flowlines.user_id"
            )
            assert attrs["flowlines.session_id"] == "sess-abc", (
                f"span '{span.name}' missing flowlines.session_id"
            )

        # Verify the leaf span has both gen_ai and flowlines attributes
        leaf_span = next(s for s in spans if s.name == "leaf-span")
        leaf_attrs = dict(leaf_span.attributes or {})
        assert leaf_attrs["gen_ai.system"] == "openai"
        assert leaf_attrs["flowlines.user_id"] == "user-42"
        assert leaf_attrs["flowlines.session_id"] == "sess-abc"


class TestFlowlinesSpanProcessor:
    def test_delegates_on_end(self) -> None:
        """FlowlinesSpanProcessor delegates on_end to inner processor."""
        inner = MagicMock()
        processor = FlowlinesSpanProcessor(inner)
        mock_span = MagicMock()

        processor.on_end(mock_span)

        inner.on_end.assert_called_once_with(mock_span)

    def test_delegates_shutdown(self) -> None:
        """FlowlinesSpanProcessor delegates shutdown to inner processor."""
        inner = MagicMock()
        processor = FlowlinesSpanProcessor(inner)

        processor.shutdown()

        inner.shutdown.assert_called_once()

    def test_delegates_force_flush(self) -> None:
        """FlowlinesSpanProcessor delegates force_flush to inner processor."""
        inner = MagicMock()
        processor = FlowlinesSpanProcessor(inner)

        processor.force_flush(timeout_millis=5000)

        inner.force_flush.assert_called_once_with(5000)
