"""Tests for tracing context managers and context propagation."""

from __future__ import annotations

import json
import time
from unittest.mock import MagicMock, patch

import pytest

from llm_observatory.tracing import (
    _current_span,
    _parse_traceparent,
    emit_completed_span,
    get_current_span_id,
    get_current_trace_id,
    get_current_traceparent,
    observe,
    span,
    trace,
)


@pytest.fixture(autouse=True)
def _reset_context():
    """Ensure each test starts with no active span context."""
    token = _current_span.set(None)
    yield
    _current_span.reset(token)


@pytest.fixture
def mock_transport():
    """Patch get_transport to capture enqueued events."""
    transport = MagicMock()
    events: list[dict] = []
    transport.enqueue.side_effect = lambda e: events.append(e)
    with patch("llm_observatory.tracing.get_transport", return_value=transport):
        yield events


@pytest.fixture
def enabled_config():
    with patch("llm_observatory.tracing.get_config", return_value={"enabled": True}):
        yield


# --- Context manager basics ---


def test_trace_sets_and_resets_context(mock_transport, enabled_config):
    assert get_current_trace_id() is None

    with trace(name="my-trace"):
        tid = get_current_trace_id()
        assert tid is not None
        sid = get_current_span_id()
        assert sid is not None

    assert get_current_trace_id() is None
    assert get_current_span_id() is None


def test_trace_sends_event_on_exit(mock_transport, enabled_config):
    with trace(name="test-trace", input={"query": "hello"}) as t:
        t.set_output({"answer": "world"})

    assert len(mock_transport) == 1
    event = mock_transport[0]
    assert event["span_type"] == "trace"
    assert event["span_name"] == "test-trace"
    assert event["trace_id"] is not None
    assert event["parent_id"] is None
    assert event["span_status"] == "ok"
    assert json.loads(event["span_input"]) == {"query": "hello"}
    assert json.loads(event["span_output"]) == {"answer": "world"}


def test_span_sends_event_on_exit(mock_transport, enabled_config):
    with trace(name="parent"):
        trace_id = get_current_trace_id()
        parent_span_id = get_current_span_id()

        with span(name="child-span") as s:
            s.set_output("done")

    # span event + trace event
    assert len(mock_transport) == 2
    span_event = mock_transport[0]
    assert span_event["span_type"] == "span"
    assert span_event["span_name"] == "child-span"
    assert span_event["trace_id"] == trace_id
    assert span_event["parent_id"] == parent_span_id


def test_nested_spans_correct_parent_hierarchy(mock_transport, enabled_config):
    with trace(name="root"):
        root_span_id = get_current_span_id()

        with span(name="level-1"):
            level1_span_id = get_current_span_id()

            with span(name="level-2"):
                pass

    # 3 events: level-2 span, level-1 span, trace
    assert len(mock_transport) == 3
    level2_event = mock_transport[0]
    level1_event = mock_transport[1]

    assert level2_event["parent_id"] == level1_span_id
    assert level1_event["parent_id"] == root_span_id


def test_generation_picks_up_trace_context(mock_transport, enabled_config):
    """build_event should add trace_id and parent_id from active context."""
    from llm_observatory.event import build_event

    with trace(name="gen-trace"):
        tid = get_current_trace_id()
        sid = get_current_span_id()

        event = build_event(
            request_kwargs={"model": "gpt-4o", "messages": []},
            response=None,
            start_time=time.time(),
            end_time=time.time() + 1,
        )

    assert event["trace_id"] == tid
    assert event["parent_id"] == sid
    assert event["span_type"] == "generation"


def test_exception_sets_error_status(mock_transport, enabled_config):
    with pytest.raises(ValueError, match="boom"):
        with trace(name="error-trace"):
            raise ValueError("boom")

    assert len(mock_transport) == 1
    assert mock_transport[0]["span_status"] == "error"


def test_exception_not_suppressed(mock_transport, enabled_config):
    with pytest.raises(RuntimeError):
        with span(name="failing-span"):
            raise RuntimeError("should propagate")


def test_disabled_config_no_events(mock_transport):
    with patch("llm_observatory.tracing.get_config", return_value={"enabled": False}):
        with trace(name="disabled-trace") as t:
            t.set_output("nope")

    assert len(mock_transport) == 0


def test_sdk_failure_does_not_crash_user_code(enabled_config):
    """Even if transport throws, user code should not be affected."""
    transport = MagicMock()
    transport.enqueue.side_effect = RuntimeError("transport broken")

    with patch("llm_observatory.tracing.get_transport", return_value=transport):
        result = None
        with trace(name="safe-trace"):
            result = 42

    assert result == 42


def test_span_outside_trace_creates_implicit_trace_id(mock_transport, enabled_config):
    with span(name="orphan-span"):
        tid = get_current_trace_id()
        assert tid is not None

    assert len(mock_transport) == 1
    event = mock_transport[0]
    assert event["trace_id"] == tid
    assert event["parent_id"] is None


# --- Cross-service propagation ---


def test_get_current_traceparent_format(mock_transport, enabled_config):
    assert get_current_traceparent() is None

    with trace(name="tp-trace"):
        tp = get_current_traceparent()
        parts = tp.split("-")
        assert len(parts) == 4
        assert parts[0] == "00"
        assert parts[3] == "01"
        # Traceparent strips hyphens; IDs are hyphenated UUIDs internally
        assert parts[1] == get_current_trace_id().replace("-", "")
        assert parts[2] == get_current_span_id().replace("-", "")


def test_get_current_traceparent_none_outside():
    assert get_current_traceparent() is None


def test_emit_completed_span_sends_event(mock_transport, enabled_config):
    now_ns = time.time_ns()
    later_ns = now_ns + 500_000_000  # 500ms later
    traceparent = "00-abcdef12345678900a1b2c3d4e5f6789-12345678abcdef0012345678abcdef00-01"

    emit_completed_span(
        name="sqs_queue_wait",
        traceparent=traceparent,
        start_ns=now_ns,
        end_ns=later_ns,
    )

    assert len(mock_transport) == 1
    event = mock_transport[0]
    assert event["span_type"] == "span"
    assert event["span_name"] == "sqs_queue_wait"
    assert event["trace_id"] == "abcdef12-3456-7890-0a1b-2c3d4e5f6789"
    assert event["parent_id"] == "12345678-abcd-ef00-1234-5678abcdef00"
    assert event["metrics"]["latency_ms"] == 500


def test_trace_with_traceparent_restores_context(mock_transport, enabled_config):
    traceparent = "00-abcdef12345678900a1b2c3d4e5f6789-12345678abcdef0012345678abcdef00-01"

    with trace(name="worker-trace", traceparent=traceparent):
        assert get_current_trace_id() == "abcdef12-3456-7890-0a1b-2c3d4e5f6789"

    assert len(mock_transport) == 1
    event = mock_transport[0]
    assert event["trace_id"] == "abcdef12-3456-7890-0a1b-2c3d4e5f6789"
    assert event["parent_id"] == "12345678-abcd-ef00-1234-5678abcdef00"
    # With a traceparent, this is a child — emitted as span, not trace
    assert event["span_type"] == "span"


def test_trace_without_traceparent_is_root(mock_transport, enabled_config):
    with trace(name="root-trace"):
        pass

    assert len(mock_transport) == 1
    assert mock_transport[0]["span_type"] == "trace"
    assert mock_transport[0]["parent_id"] is None


# --- Update method ---


def test_trace_update_name(mock_transport, enabled_config):
    with trace(name="old-name") as t:
        t.update(name="new-name")

    assert mock_transport[0]["span_name"] == "new-name"


# --- Traceparent parser ---


def test_parse_traceparent_valid():
    result = _parse_traceparent("00-abcdef12345678900a1b2c3d4e5f6789-12345678abcdef0012345678abcdef00-01")
    assert result == ("abcdef12-3456-7890-0a1b-2c3d4e5f6789", "12345678-abcd-ef00-1234-5678abcdef00")


def test_parse_traceparent_invalid():
    assert _parse_traceparent("invalid") is None
    assert _parse_traceparent("") is None
    assert _parse_traceparent("00-not_a_valid_hex_string_here!!-1234567890abcdef1234567890abcdef-01") is None
    assert _parse_traceparent("00-abcdef12345678900a1b2c3d4e5f6789-ZZZZ_not_valid_hex_string!!!!-01") is None


def test_emit_completed_span_invalid_traceparent(mock_transport, enabled_config):
    emit_completed_span(
        name="bad-span",
        traceparent="garbage",
        start_ns=1000,
        end_ns=2000,
    )
    assert len(mock_transport) == 0


# --- @observe decorator ---


def test_observe_creates_trace_at_top_level(mock_transport, enabled_config):
    @observe
    def my_pipeline(query: str):
        return {"answer": "42"}

    result = my_pipeline("hello")
    assert result == {"answer": "42"}
    assert len(mock_transport) == 1
    event = mock_transport[0]
    assert event["span_type"] == "trace"
    assert event["span_name"] == "my_pipeline"
    assert event["span_status"] == "ok"


def test_observe_captures_input_and_output(mock_transport, enabled_config):
    @observe
    def search(query: str, top_k: int = 5):
        return ["doc1", "doc2"]

    search("what is X?", top_k=3)
    event = mock_transport[0]
    assert json.loads(event["span_input"]) == {"query": "what is X?", "top_k": 3}
    assert json.loads(event["span_output"]) == ["doc1", "doc2"]


def test_observe_nested_becomes_span(mock_transport, enabled_config):
    @observe
    def inner():
        return "done"

    @observe
    def outer():
        return inner()

    outer()
    assert len(mock_transport) == 2
    inner_event = mock_transport[0]
    outer_event = mock_transport[1]
    assert inner_event["span_type"] == "span"
    assert outer_event["span_type"] == "trace"
    assert inner_event["trace_id"] == outer_event["trace_id"]
    assert inner_event["parent_id"] is not None


def test_observe_with_custom_name(mock_transport, enabled_config):
    @observe(name="custom-step")
    def step():
        pass

    step()
    assert mock_transport[0]["span_name"] == "custom-step"


def test_observe_capture_input_false(mock_transport, enabled_config):
    @observe(capture_input=False)
    def secret_fn(api_key: str):
        return "ok"

    secret_fn("sk-secret")
    event = mock_transport[0]
    assert "span_input" not in event


def test_observe_capture_output_false(mock_transport, enabled_config):
    @observe(capture_output=False)
    def fn():
        return {"sensitive": "data"}

    fn()
    event = mock_transport[0]
    assert "span_output" not in event


def test_observe_exception_sets_error(mock_transport, enabled_config):
    @observe
    def failing():
        raise ValueError("boom")

    with pytest.raises(ValueError, match="boom"):
        failing()

    assert len(mock_transport) == 1
    assert mock_transport[0]["span_status"] == "error"


def test_observe_async_function(mock_transport, enabled_config):
    import asyncio

    @observe
    async def async_fn(x: int):
        return x * 2

    result = asyncio.run(async_fn(21))
    assert result == 42
    assert len(mock_transport) == 1
    assert mock_transport[0]["span_name"] == "async_fn"
    assert json.loads(mock_transport[0]["span_output"]) == 42


def test_observe_mixed_with_context_manager(mock_transport, enabled_config):
    """@observe and with span() compose within the same trace."""

    @observe
    def pipeline():
        with span(name="manual-step") as s:
            s.set_output("step-done")
        return "pipeline-done"

    pipeline()
    assert len(mock_transport) == 2
    span_event = mock_transport[0]
    trace_event = mock_transport[1]
    assert span_event["span_type"] == "span"
    assert span_event["span_name"] == "manual-step"
    assert trace_event["span_type"] == "trace"
    assert span_event["trace_id"] == trace_event["trace_id"]


def test_observe_preserves_function_metadata():
    @observe
    def documented_fn():
        """My docstring."""
        pass

    assert documented_fn.__name__ == "documented_fn"
    assert documented_fn.__doc__ == "My docstring."


# --- flush() ---


def test_flush_drains_pending_events(mock_transport, enabled_config):
    """flush() should synchronously send all pending events."""
    import llm_observatory

    with trace(name="flush-test"):
        pass

    assert len(mock_transport) == 1

    # flush() should not raise even when transport is already drained
    llm_observatory.flush()
