import asyncio
import pytest

from bandito.context import (
    SpanContext,
    TraceContext,
    get_current_span,
    get_current_trace,
    spawn_task,
)


class TestCurrentContextAccessors:
    """Tests for get_current_trace() and get_current_span()."""

    def test_get_current_trace_none_outside_context(self):
        assert get_current_trace() is None

    def test_get_current_span_none_outside_context(self):
        assert get_current_span() is None

    def test_get_current_trace_inside_trace_context(self):
        with TraceContext("proj") as t:
            assert get_current_trace() is t

    def test_get_current_span_inside_span_context(self):
        with TraceContext("proj") as t:
            with t.span("gen", "llm") as s:
                assert get_current_span() is s

    def test_get_current_span_none_after_span_exits(self):
        with TraceContext("proj") as t:
            with t.span("gen", "llm") as s:
                pass
            assert get_current_span() is None
        assert get_current_trace() is None


class TestSpawnTask:
    """Tests for spawn_task() context propagation."""

    def test_spawn_task_inherits_trace_context(self):
        """Child task should see parent trace context via spawn_task."""
        async def main():
            results = {}

            async def worker():
                results["trace"] = get_current_trace()
                results["span"] = get_current_span()

            with TraceContext("proj") as t:
                with t.span("parent", "llm") as parent:
                    task = spawn_task(worker())
                    await task

            # Verify the task captured the context
            assert results["trace"] is not None
            assert results["trace"].trace.project == "proj"
            assert results["span"] is not None
            assert results["span"].span.name == "parent"

        asyncio.run(main())

    def test_create_task_copies_context(self):
        """Python 3.12+ copies context by default for create_task."""
        async def main():
            results = {}

            async def worker():
                results["trace"] = get_current_trace()

            with TraceContext("proj") as t:
                with t.span("parent", "llm"):
                    # create_task in Python 3.12+ copies context
                    task = asyncio.create_task(worker())
                    await task

            # Context IS preserved (Python 3.12+ behavior)
            assert results["trace"] is not None
            assert results["trace"].trace.project == "proj"

        asyncio.run(main())

    def test_spawn_task_preserves_context_across_tasks(self):
        """Verify spawn_task preserves context even when parent exits."""
        async def main():
            captured_contexts = []

            async def inner_worker():
                captured_contexts.append((get_current_trace(), get_current_span()))

            with TraceContext("proj") as t:
                with t.span("parent", "llm") as s:
                    # Task will run after parent span exits but within trace
                    task = spawn_task(inner_worker())
                    await task
                    # After task completes, parent span still active
                    assert get_current_span() is s

            assert captured_contexts[0][0] is not None

        asyncio.run(main())


class TestTraceContext:
    def test_basic_lifecycle(self):
        traces = []

        class Collector:
            def write(self, trace):
                traces.append(trace)

        with TraceContext("proj", writer=Collector(), input="hi") as t:
            with t.span("gen", "llm") as s:
                s.done(output="hello")
            t.done(output="hello")

        assert len(traces) == 1
        trace = traces[0]
        assert trace.project == "proj"
        assert trace.input == "hi"
        assert trace.output == "hello"
        assert trace.status == "ok"
        assert trace.ended_at is not None
        assert len(trace.spans) == 1
        assert trace.spans[0].name == "gen"
        assert trace.spans[0].output == "hello"

    def test_auto_nesting(self):
        traces = []

        class Collector:
            def write(self, trace):
                traces.append(trace)

        with TraceContext("proj", writer=Collector()) as t:
            with t.span("parent", "agent") as parent:
                with t.span("child", "llm") as child:
                    child.done(output="nested")
                parent.done(output="done")

        trace = traces[0]
        assert len(trace.spans) == 2
        parent_span = next(s for s in trace.spans if s.name == "parent")
        child_span = next(s for s in trace.spans if s.name == "child")
        assert child_span.parent_id == parent_span.span_id
        assert parent_span.parent_id is None

    def test_error_capture_on_trace(self):
        traces = []

        class Collector:
            def write(self, trace):
                traces.append(trace)

        with pytest.raises(ValueError, match="boom"):
            with TraceContext("proj", writer=Collector()) as t:
                with t.span("gen", "llm") as s:
                    s.done(output="ok")
                raise ValueError("boom")

        trace = traces[0]
        assert trace.status == "error"
        assert "ValueError: boom" in trace.error

    def test_error_capture_on_span(self):
        traces = []

        class Collector:
            def write(self, trace):
                traces.append(trace)

        with pytest.raises(RuntimeError, match="fail"):
            with TraceContext("proj", writer=Collector()) as t:
                with t.span("gen", "llm") as s:
                    raise RuntimeError("fail")

        trace = traces[0]
        # Both span and trace should capture error
        assert trace.spans[0].status == "error"
        assert "RuntimeError: fail" in trace.spans[0].error
        assert trace.status == "error"

    def test_timing(self):
        traces = []

        class Collector:
            def write(self, trace):
                traces.append(trace)

        with TraceContext("proj", writer=Collector()) as t:
            with t.span("gen", "llm") as s:
                s.done()

        trace = traces[0]
        assert trace.started_at <= trace.ended_at
        span = trace.spans[0]
        assert span.started_at <= span.ended_at

    def test_no_writer(self):
        """TraceContext with no writer should not raise."""
        with TraceContext("proj") as t:
            with t.span("gen", "llm") as s:
                s.done(output="ok")
            t.done(output="ok")
        # Just assert it completed without error
        assert t.trace.output == "ok"

    def test_span_kwargs(self):
        traces = []

        class Collector:
            def write(self, trace):
                traces.append(trace)

        with TraceContext("proj", writer=Collector()) as t:
            with t.span("gen", "llm", model="gpt-4o", provider="openai") as s:
                s.done(output="hi", input_tokens=10, output_tokens=5)

        span = traces[0].spans[0]
        assert span.model == "gpt-4o"
        assert span.provider == "openai"
        assert span.input_tokens == 10
        assert span.output_tokens == 5

    def test_deep_nesting(self):
        traces = []

        class Collector:
            def write(self, trace):
                traces.append(trace)

        with TraceContext("proj", writer=Collector()) as t:
            with t.span("agent", "agent") as a:
                with t.span("llm1", "llm") as l1:
                    l1.done()
                with t.span("tool", "tool") as tool:
                    with t.span("llm2", "llm") as l2:
                        l2.done()
                    tool.done()
                a.done()

        trace = traces[0]
        assert len(trace.spans) == 4
        spans = {s.name: s for s in trace.spans}
        assert spans["agent"].parent_id is None
        assert spans["llm1"].parent_id == spans["agent"].span_id
        assert spans["tool"].parent_id == spans["agent"].span_id
        assert spans["llm2"].parent_id == spans["tool"].span_id

    def test_done_with_mapper_skips_duplicate_llm_spans(self):
        """When explicit LLM spans exist, mapper LLM spans are skipped."""
        traces = []

        class Collector:
            def write(self, trace):
                traces.append(trace)

        # Simulate: user creates explicit LLM span, then passes a framework
        # result to t.done() — mapper should NOT add another LLM span
        with TraceContext("proj", writer=Collector()) as t:
            with t.span("gen", "llm", model="minimax/minimax-m2.5", provider="openrouter") as s:
                s.done(output="hello", input_tokens=100, output_tokens=50)
            # Simulate mapper extraction by monkey-patching
            from unittest.mock import patch
            from bandito.models import Span

            fake_llm = Span(name="minimax/minimax-m2.5", kind="llm", model="minimax/minimax-m2.5")
            fake_tool = Span(name="get_weather", kind="tool", input={"city": "NYC"})
            with patch("bandito.mappers.extract", return_value=("hello", [fake_llm, fake_tool])):
                t.done("fake_result")

        trace = traces[0]
        llm_spans = [s for s in trace.spans if s.kind == "llm"]
        tool_spans = [s for s in trace.spans if s.kind == "tool"]
        # Only the explicit LLM span, not the mapper-extracted one
        assert len(llm_spans) == 1
        assert llm_spans[0].name == "gen"
        # Tool spans from mapper should still be added
        assert len(tool_spans) == 1
        assert tool_spans[0].name == "get_weather"

    def test_done_with_mapper_adds_spans_when_no_explicit_llm(self):
        """When no explicit LLM spans exist, mapper spans are all added."""
        traces = []

        class Collector:
            def write(self, trace):
                traces.append(trace)

        with TraceContext("proj", writer=Collector()) as t:
            from unittest.mock import patch
            from bandito.models import Span

            fake_llm = Span(name="gpt-4o", kind="llm", model="gpt-4o")
            fake_tool = Span(name="search", kind="tool")
            with patch("bandito.mappers.extract", return_value=("result", [fake_llm, fake_tool])):
                t.done("fake_result")

        trace = traces[0]
        assert len(trace.spans) == 2
        assert trace.spans[0].name == "gpt-4o"
        assert trace.spans[1].name == "search"
