"""Tests for the manual tracing API: trace(), span(), outcome()."""
from __future__ import annotations

import pytest

from fibinaci.tracing.context import (
    get_current_span_id,
    get_current_trace_id,
    outcome,
    span,
    trace,
)
from fibinaci.tracing.span_builder import SpanData, estimate_cost
from fibinaci.tracing.decorators import observe
from fibinaci._testing.mock_client import MockFibinaci


class TestTraceContext:
    def test_trace_sets_trace_id(self, mock_fibinaci: MockFibinaci):
        with trace("test-workflow") as tid:
            assert tid is not None
            assert len(tid) == 32
            assert get_current_trace_id() == tid

    def test_trace_id_cleared_after_exit(self, mock_fibinaci: MockFibinaci):
        with trace("t"):
            pass
        assert get_current_trace_id() is None

    def test_nested_spans(self, mock_fibinaci: MockFibinaci):
        with trace("t") as tid:
            with span("parent", span_kind="agent_planning") as parent:
                parent_id = parent.span_id
                with span("child", span_kind="model_call") as child:
                    assert child.parent_span_id == parent_id

        spans = mock_fibinaci.get_spans()
        assert len(spans) == 2

    def test_span_captures_timing(self, mock_fibinaci: MockFibinaci):
        with trace("t"):
            with span("s", span_kind="model_call") as sd:
                import time
                time.sleep(0.01)

        spans = mock_fibinaci.get_spans()
        assert len(spans) == 1
        assert spans[0]["duration_ms"] >= 10

    def test_span_captures_error(self, mock_fibinaci: MockFibinaci):
        with trace("t"):
            with pytest.raises(ValueError):
                with span("s", span_kind="tool_call") as sd:
                    raise ValueError("broken")

        spans = mock_fibinaci.get_spans()
        assert spans[0]["status"] == "error"
        assert spans[0]["error_type"] == "ValueError"
        assert spans[0]["error_message"] == "broken"

    def test_standalone_span_outside_trace(self, mock_fibinaci: MockFibinaci):
        """span() without trace() should still capture data."""
        with span("standalone", span_kind="custom") as sd:
            assert sd.trace_id  # auto-generated

    def test_outcome_captured(self, mock_fibinaci: MockFibinaci):
        with trace("t"):
            with span("s", span_kind="model_call"):
                pass
            outcome(success=True, score=0.95)

        outcomes = mock_fibinaci.get_outcomes()
        assert len(outcomes) == 1
        assert outcomes[0]["success"] is True
        assert outcomes[0]["score"] == 0.95

    def test_trace_propagates_session_id(self, mock_fibinaci: MockFibinaci):
        with trace("t", session_id="sess_123"):
            with span("s", span_kind="model_call"):
                pass

        spans = mock_fibinaci.get_spans()
        assert spans[0].get("session_id") == "sess_123"


class TestSpanBuilder:
    def test_to_dict_minimal(self):
        sd = SpanData(trace_id="a" * 32, span_id="b" * 16)
        d = sd.to_dict()
        assert d["trace_id"] == "a" * 32
        assert d["status"] == "ok"

    def test_set_tokens_auto_calculates_total(self):
        sd = SpanData()
        sd.set_tokens(input=100, output=50)
        assert sd.tokens_total == 150

    def test_set_model_call_auto_calculates_cost(self):
        sd = SpanData(tokens_input=1_000_000, tokens_output=500_000)
        sd.set_model_call(model_name="gpt-4o")
        assert sd.cost_usd is not None
        assert float(sd.cost_usd) > 0

    def test_to_dict_includes_extensions(self):
        sd = SpanData()
        sd.set_model_call(model_name="gpt-4o", temperature=0.7)
        sd.set_tool_call(tool_name="search")
        d = sd.to_dict()
        assert d["model_call"]["model_name"] == "gpt-4o"
        assert d["tool_call"]["tool_name"] == "search"


class TestEstimateCost:
    def test_known_model(self):
        cost = estimate_cost("gpt-4o", 1_000_000, 1_000_000)
        assert cost is not None
        assert float(cost) == pytest.approx(12.50, abs=0.01)  # 2.5 + 10.0

    def test_unknown_model_returns_none(self):
        assert estimate_cost("unknown-model", 1000, 1000) is None

    def test_no_tokens_returns_none(self):
        assert estimate_cost("gpt-4o", None, None) is None

    def test_no_model_returns_none(self):
        assert estimate_cost(None, 1000, 1000) is None


class TestObserveDecorator:
    def test_sync_function(self, mock_fibinaci: MockFibinaci):
        @observe("my-func", span_kind="tool_call")
        def my_func():
            return 42

        with trace("t"):
            result = my_func()

        assert result == 42
        spans = mock_fibinaci.get_spans()
        assert any(s["span_kind"] == "tool_call" for s in spans)

    @pytest.mark.asyncio
    async def test_async_function(self, mock_fibinaci: MockFibinaci):
        @observe("async-func", span_kind="model_call")
        async def async_func():
            return "hello"

        with trace("t"):
            result = await async_func()

        assert result == "hello"
        spans = mock_fibinaci.get_spans()
        assert any(s["span_kind"] == "model_call" for s in spans)

    def test_preserves_exception(self, mock_fibinaci: MockFibinaci):
        @observe()
        def failing():
            raise RuntimeError("fail")

        with trace("t"):
            with pytest.raises(RuntimeError):
                failing()

        spans = mock_fibinaci.get_spans()
        assert spans[0]["status"] == "error"
