from agenx._internal import (
    make_span_id,
    make_trace_id,
    now,
    get_current_span,
    push_span,
    pop_span
)
from agenx.trace.spans.span import LiveSpan
from agenx.trace.spans.models import (
    Span,
    SpanType
)

def make_live(callbacks=None, **overrides) -> LiveSpan:
    return LiveSpan(generate_span(**overrides), on_end_callbacks=callbacks or [])


def generate_span(**overrides) -> Span:
    defaults = dict(
        span_id=make_span_id(),
        trace_id=make_trace_id(),
        parent_id=None,
        name="test_span",
        type=SpanType.STEP,
        start_time=now(),
    )

    defaults.update(**overrides)
    return Span(**defaults)


def test_default_span_is_none():
    """Before any span is pushed, context must be clean"""
    assert get_current_span() is None


def test_push_makes_current():
    span = make_live()
    token = push_span(span)
    try:
        assert get_current_span() is span
    finally:
        pop_span(token)


def test_nested_restores_outer():
    outer, inner = make_live("outer"), make_live("inner")
    t1 = push_span(outer)
    t2 = push_span(inner)
    
    assert get_current_span() is inner
    pop_span(t2)
    
    assert get_current_span() is outer  
    pop_span(t1)

    assert get_current_span() is None


def test_pop_runs_even_on_exception():
    span = make_live()
    token = push_span(span)
    try:
        raise ValueError("body raised")
    except ValueError:
        pass
    finally:
        pop_span(token)
    assert get_current_span() is None
