from __future__ import annotations

import pytest
from opentelemetry import context as otel_context
from opentelemetry import trace
from opentelemetry.trace import (
    NonRecordingSpan,
    SpanContext,
    TraceFlags,
    TraceState,
    set_span_in_context,
)

from terminaluse.lib.sdk.fastacp.base.base_acp_server import BaseACPServer

TRACEPARENT = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"


@pytest.mark.unit
def test_extract_trace_context_headers_filters_and_normalizes() -> None:
    headers = {
        "TraceParent": TRACEPARENT,
        "TrAcEsTaTe": "rojo=00f067aa0ba902b7",
        "Baggage": "team=platform",
        "X-Request-Id": "req-123",
        "Authorization": "blocked",
    }

    assert BaseACPServer._extract_trace_context_headers(headers) == {
        "traceparent": TRACEPARENT,
        "tracestate": "rojo=00f067aa0ba902b7",
        "baggage": "team=platform",
    }


@pytest.mark.unit
def test_resolve_background_context_uses_extracted_context_without_current_span() -> None:
    # Reset active context for deterministic assertions.
    reset_token = otel_context.attach(otel_context.Context())
    try:
        ctx = BaseACPServer._resolve_background_context({"traceparent": TRACEPARENT})
    finally:
        otel_context.detach(reset_token)

    span_context = ctx.run(trace.get_current_span).get_span_context()
    assert span_context.is_valid
    assert span_context.is_remote
    assert format(span_context.trace_id, "032x") == "4bf92f3577b34da6a3ce929d0e0e4736"
    assert format(span_context.span_id, "016x") == "00f067aa0ba902b7"


@pytest.mark.unit
def test_resolve_background_context_prefers_existing_current_span() -> None:
    current_span_context = SpanContext(
        trace_id=int("11111111111111111111111111111111", 16),
        span_id=int("2222222222222222", 16),
        is_remote=False,
        trace_flags=TraceFlags(TraceFlags.SAMPLED),
        trace_state=TraceState(),
    )

    active_context = set_span_in_context(NonRecordingSpan(current_span_context))
    token = otel_context.attach(active_context)
    try:
        ctx = BaseACPServer._resolve_background_context({"traceparent": TRACEPARENT})
    finally:
        otel_context.detach(token)

    span_context = ctx.run(trace.get_current_span).get_span_context()
    assert span_context.is_valid
    assert not span_context.is_remote
    assert format(span_context.trace_id, "032x") == "11111111111111111111111111111111"
    assert format(span_context.span_id, "016x") == "2222222222222222"
