import logging
import os
import subprocess
import sys
import textwrap
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any, cast
from unittest.mock import MagicMock, patch

import httpx
import pytest

from agents.tracing import flush_traces, get_trace_provider
from agents.tracing.processor_interface import TracingExporter, TracingProcessor
from agents.tracing.processors import BackendSpanExporter, BatchTraceProcessor
from agents.tracing.provider import DefaultTraceProvider, TraceProvider
from agents.tracing.span_data import AgentSpanData
from agents.tracing.spans import Span, SpanImpl
from agents.tracing.traces import Trace, TraceImpl


def get_span(processor: TracingProcessor) -> SpanImpl[AgentSpanData]:
    """Create a minimal agent span for testing processors."""
    return SpanImpl(
        trace_id="test_trace_id",
        span_id="test_span_id",
        parent_id=None,
        processor=processor,
        span_data=AgentSpanData(name="test_agent"),
        tracing_api_key=None,
    )


def get_trace(processor: TracingProcessor) -> TraceImpl:
    """Create a minimal trace."""
    return TraceImpl(
        name="test_trace",
        trace_id="test_trace_id",
        group_id="test_session_id",
        metadata={},
        processor=processor,
        tracing_api_key=None,
    )


@pytest.fixture
def mocked_exporter():
    exporter = MagicMock()
    exporter.export = MagicMock()
    return exporter


def test_batch_trace_processor_on_trace_start(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=0.1)
    test_trace = get_trace(processor)

    processor.on_trace_start(test_trace)
    assert processor._queue.qsize() == 1, "Trace should be added to the queue"

    # Shutdown to clean up the worker thread
    processor.shutdown()


def test_batch_trace_processor_on_span_end(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=0.1)
    test_span = get_span(processor)

    processor.on_span_end(test_span)
    assert processor._queue.qsize() == 1, "Span should be added to the queue"

    # Shutdown to clean up the worker thread
    processor.shutdown()


def test_batch_trace_processor_queue_full(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, max_queue_size=2, schedule_delay=0.1)
    # Fill the queue
    processor.on_trace_start(get_trace(processor))
    processor.on_trace_start(get_trace(processor))
    assert processor._queue.full() is True

    # Next item should not be queued
    processor.on_trace_start(get_trace(processor))
    assert processor._queue.qsize() == 2, "Queue should not exceed max_queue_size"

    processor.on_span_end(get_span(processor))
    assert processor._queue.qsize() == 2, "Queue should not exceed max_queue_size"

    processor.shutdown()


def test_batch_processor_doesnt_enqueue_on_trace_end_or_span_start(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter)

    processor.on_trace_start(get_trace(processor))
    assert processor._queue.qsize() == 1, "Trace should be queued"

    processor.on_span_start(get_span(processor))
    assert processor._queue.qsize() == 1, "Span should not be queued"

    processor.on_span_end(get_span(processor))
    assert processor._queue.qsize() == 2, "Span should be queued"

    processor.on_trace_end(get_trace(processor))
    assert processor._queue.qsize() == 2, "Nothing new should be queued"

    processor.shutdown()


def test_batch_trace_processor_force_flush(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, max_batch_size=2, schedule_delay=5.0)

    processor.on_trace_start(get_trace(processor))
    processor.on_span_end(get_span(processor))
    processor.on_span_end(get_span(processor))

    processor.force_flush()

    # Ensure exporter.export was called with all items
    # Because max_batch_size=2, it may have been called multiple times
    total_exported = 0
    for call_args in mocked_exporter.export.call_args_list:
        batch = call_args[0][0]  # first positional arg to export() is the items list
        total_exported += len(batch)

    # We pushed 3 items; ensure they all got exported
    assert total_exported == 3

    processor.shutdown()


def test_batch_trace_processor_force_flush_waits_for_in_flight_background_export():
    export_started = threading.Event()
    export_continue = threading.Event()

    class BlockingExporter(TracingExporter):
        def export(self, items: list[Trace | Span[Any]]) -> None:
            export_started.set()
            assert export_continue.wait(timeout=2.0)

    processor = BatchTraceProcessor(exporter=BlockingExporter(), schedule_delay=0.01)
    processor.on_trace_start(get_trace(processor))

    assert export_started.wait(timeout=2.0)

    flush_thread = threading.Thread(target=processor.force_flush)
    flush_thread.start()

    time.sleep(0.1)
    assert flush_thread.is_alive(), "force_flush() should wait for an in-flight export"

    export_continue.set()
    flush_thread.join(timeout=2.0)

    assert not flush_thread.is_alive()

    processor.shutdown()


def test_batch_trace_processor_shutdown_flushes(mocked_exporter):
    processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=5.0)
    processor.on_trace_start(get_trace(processor))
    processor.on_span_end(get_span(processor))
    qsize_before = processor._queue.qsize()
    assert qsize_before == 2

    processor.shutdown()

    # Ensure everything was exported after shutdown
    total_exported = 0
    for call_args in mocked_exporter.export.call_args_list:
        batch = call_args[0][0]
        total_exported += len(batch)

    assert total_exported == 2, "All items in the queue should be exported upon shutdown"


def test_batch_trace_processor_shutdown_timeout_returns_when_exporter_blocks(
    caplog: pytest.LogCaptureFixture,
) -> None:
    export_started = threading.Event()
    release_export = threading.Event()

    class BlockingExporter(TracingExporter):
        def export(self, items: list[Trace | Span[Any]]) -> None:
            export_started.set()
            release_export.wait(timeout=5.0)

    processor = BatchTraceProcessor(
        exporter=BlockingExporter(),
        max_queue_size=1,
        schedule_delay=60.0,
        export_trigger_ratio=1.0,
    )
    processor.on_span_end(get_span(processor))

    assert export_started.wait(timeout=2.0)

    start = time.monotonic()
    with caplog.at_level(logging.WARNING):
        processor.shutdown(timeout=0.05)
    elapsed = time.monotonic() - start

    assert elapsed < 0.5
    assert "shutdown timeout reached" in caplog.text

    release_export.set()
    if processor._worker_thread:
        processor._worker_thread.join(timeout=2.0)


def test_batch_trace_processor_shutdown_passes_deadline_to_exporter() -> None:
    seen_deadlines: list[float | None] = []

    class DeadlineExporter(TracingExporter):
        def export(self, items: list[Trace | Span[Any]]) -> None:
            raise AssertionError("shutdown should use the deadline-aware exporter path")

        def _export_with_deadline(
            self, items: list[Trace | Span[Any]], deadline: float | None
        ) -> None:
            seen_deadlines.append(deadline)

    processor = BatchTraceProcessor(exporter=DeadlineExporter())
    processor._queue.put_nowait(get_span(processor))

    processor.shutdown(timeout=1.0)

    assert len(seen_deadlines) == 1
    assert seen_deadlines[0] is not None


def test_batch_trace_processor_survives_exporter_exception():
    """A failing exporter must not kill the background worker thread.

    Previously, an exception raised inside ``exporter.export`` propagated out of
    ``_export_batches`` and killed the ``_run`` thread, causing all subsequent
    spans to silently accumulate in the queue until it filled up.
    """

    class FlakyExporter(TracingExporter):
        def __init__(self) -> None:
            self.call_count = 0
            self.exported: list[Trace | Span[Any]] = []

        def export(self, items: list[Trace | Span[Any]]) -> None:
            self.call_count += 1
            if self.call_count == 1:
                raise RuntimeError("simulated exporter failure")
            self.exported.extend(items)

    exporter = FlakyExporter()
    processor = BatchTraceProcessor(exporter, schedule_delay=0.05, max_batch_size=1)
    processor.on_span_end(get_span(processor))
    processor.on_span_end(get_span(processor))
    processor.on_span_end(get_span(processor))

    # Give the worker time to encounter the failure and continue processing.
    time.sleep(0.3)

    assert processor._worker_thread is not None
    assert processor._worker_thread.is_alive(), "Worker thread must survive an exporter exception"

    processor.shutdown(timeout=2.0)

    # First batch raised; the remaining two items must still have been exported.
    assert len(exporter.exported) == 2
    assert exporter.call_count >= 3


def test_batch_trace_processor_scheduled_export(mocked_exporter):
    """
    Tests that items are automatically exported when the schedule_delay expires.
    We mock time.time() so we can trigger the condition without waiting in real time.
    """
    with patch("time.time") as mock_time:
        base_time = 1000.0
        mock_time.return_value = base_time

        processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=1.0)

        processor.on_span_end(get_span(processor))  # queue size = 1

        # Now artificially advance time beyond the next export time
        mock_time.return_value = base_time + 2.0  # > base_time + schedule_delay
        # Let the background thread run a bit
        time.sleep(0.3)

        # Check that exporter.export was eventually called
        # Because the background thread runs, we might need a small sleep
        processor.shutdown()

    total_exported = 0
    for call_args in mocked_exporter.export.call_args_list:
        batch = call_args[0][0]
        total_exported += len(batch)

    assert total_exported == 1, "Item should be exported after scheduled delay"


def test_flush_traces_delegates_to_default_trace_provider():
    provider = DefaultTraceProvider()
    mock_processor = MagicMock()
    provider.register_processor(mock_processor)

    with patch("agents.tracing.setup.GLOBAL_TRACE_PROVIDER", provider):
        flush_traces()

    mock_processor.force_flush.assert_called_once()


def test_flush_traces_is_importable_from_top_level_agents_package():
    from agents import flush_traces as top_level_flush_traces

    assert top_level_flush_traces is flush_traces


def test_default_trace_provider_force_flush_respects_disabled_flag():
    provider = DefaultTraceProvider()
    mock_processor = MagicMock()
    provider.register_processor(mock_processor)

    provider.set_disabled(True)
    provider.force_flush()

    mock_processor.force_flush.assert_not_called()


def test_trace_provider_force_flush_and_shutdown_default_to_noops():
    class MinimalProvider(TraceProvider):
        def register_processor(self, processor: TracingProcessor) -> None:
            pass

        def set_processors(self, processors: list[TracingProcessor]) -> None:
            pass

        def get_current_trace(self):
            return None

        def get_current_span(self):
            return None

        def set_disabled(self, disabled: bool) -> None:
            pass

        def time_iso(self) -> str:
            return ""

        def gen_trace_id(self) -> str:
            return "trace_123"

        def gen_span_id(self) -> str:
            return "span_123"

        def gen_group_id(self) -> str:
            return "group_123"

        def create_trace(
            self,
            name,
            trace_id=None,
            group_id=None,
            metadata=None,
            disabled=False,
            tracing=None,
        ):
            raise NotImplementedError

        def create_span(self, span_data, span_id=None, parent=None, disabled=False):
            raise NotImplementedError

    provider = MinimalProvider()
    provider.force_flush()
    provider.shutdown()


def test_get_trace_provider_force_flush_flushes_default_processor(mocked_exporter):
    provider = DefaultTraceProvider()
    processor = BatchTraceProcessor(exporter=mocked_exporter, schedule_delay=60.0)
    provider.register_processor(processor)

    with patch("agents.tracing.setup.GLOBAL_TRACE_PROVIDER", provider):
        processor.on_trace_start(get_trace(processor))
        processor.on_span_end(get_span(processor))

        get_trace_provider().force_flush()

    total_exported = sum(
        len(call_args[0][0]) for call_args in mocked_exporter.export.call_args_list
    )
    assert total_exported == 2
    processor.shutdown()


def mock_processor():
    processor = MagicMock()
    processor.on_trace_start = MagicMock()
    processor.on_span_end = MagicMock()
    return processor


@patch("httpx.Client")
def test_backend_span_exporter_no_items(mock_client):
    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([])
    # No calls should be made if there are no items
    mock_client.return_value.post.assert_not_called()
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_no_api_key(mock_client):
    # Ensure that os.environ is empty (sometimes devs have the openai api key set in their env)

    with patch.dict(os.environ, {}, clear=True):
        exporter = BackendSpanExporter(api_key=None)
        exporter.export([get_span(mock_processor())])

        # Should log an error and return without calling post
        mock_client.return_value.post.assert_not_called()
        exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_2xx_success(mock_client):
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([get_span(mock_processor()), get_trace(mock_processor())])

    # Should have called post exactly once
    mock_client.return_value.post.assert_called_once()
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_4xx_client_error(mock_client):
    mock_response = MagicMock()
    mock_response.status_code = 400
    mock_response.text = "Bad Request"
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([get_span(mock_processor())])

    # 4xx should not be retried
    mock_client.return_value.post.assert_called_once()
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_5xx_retry(mock_client):
    mock_response = MagicMock()
    mock_response.status_code = 500

    # Make post() return 500 every time
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key", max_retries=3, base_delay=0.1, max_delay=0.2)
    with patch.object(exporter._shutdown_event, "wait", return_value=False) as wait_for_retry:
        exporter.export([get_span(mock_processor())])

    # Should retry up to max_retries times
    assert mock_client.return_value.post.call_count == 3
    assert wait_for_retry.call_count == 2

    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_deadline_stops_during_5xx_retry_backoff(mock_client):
    mock_response = MagicMock()
    mock_response.status_code = 504
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key", max_retries=3, base_delay=1.0)
    with patch("time.sleep") as sleep_for_retry:
        exporter._export_with_deadline(
            [get_span(mock_processor())], deadline=time.monotonic() + 0.01
        )

    assert mock_client.return_value.post.call_count == 1
    sleep_for_retry.assert_called_once()
    assert sleep_for_retry.call_args.args[0] <= 0.1

    exporter.close()


@patch("httpx.Client")
def test_batch_trace_processor_shutdown_interrupts_exporter_retry_backoff(mock_client):
    post_called = threading.Event()
    mock_response = MagicMock()
    mock_response.status_code = 504

    def post(**kwargs: Any) -> Any:
        post_called.set()
        return mock_response

    mock_client.return_value.post.side_effect = post

    exporter = BackendSpanExporter(
        api_key="test_key",
        max_retries=100,
        base_delay=10.0,
        max_delay=10.0,
    )
    processor = BatchTraceProcessor(
        exporter=exporter,
        max_queue_size=1,
        max_batch_size=1,
        schedule_delay=60.0,
        export_trigger_ratio=1.0,
    )

    processor.on_span_end(get_span(processor))
    assert post_called.wait(timeout=2.0)

    start = time.monotonic()
    processor.shutdown(timeout=1.0)
    elapsed = time.monotonic() - start

    assert elapsed < 0.5
    assert processor._worker_thread is not None
    assert not processor._worker_thread.is_alive()
    assert mock_client.return_value.post.call_count == 1

    exporter.close()


@patch("httpx.Client")
def test_batch_trace_processor_shutdown_without_timeout_preserves_export_retries(mock_client):
    mock_response = MagicMock()
    mock_response.status_code = 504
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(
        api_key="test_key",
        max_retries=3,
        base_delay=0.1,
        max_delay=0.2,
    )
    processor = BatchTraceProcessor(exporter=exporter)
    processor._queue.put_nowait(get_span(processor))

    with patch.object(exporter._shutdown_event, "wait", return_value=False) as wait_for_retry:
        processor.shutdown(timeout=None)

    assert mock_client.return_value.post.call_count == 3
    assert wait_for_retry.call_count == 2

    exporter.close()


@pytest.mark.serial
def test_tracing_atexit_cleanup_timeout_preserves_process_exit_code_on_504() -> None:
    request_seen = threading.Event()

    class Always504Handler(BaseHTTPRequestHandler):
        def do_POST(self) -> None:
            request_seen.set()
            self.send_response(504)
            self.end_headers()
            self.wfile.write(b"gateway timeout")

        def log_message(self, format: str, *args: Any) -> None:
            return

    server = ThreadingHTTPServer(("127.0.0.1", 0), Always504Handler)
    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
    server_thread.start()

    script = textwrap.dedent(
        f"""
        import sys
        import time

        from agents.tracing import custom_span, trace
        from agents.tracing.processors import BackendSpanExporter, BatchTraceProcessor
        from agents.tracing.provider import DefaultTraceProvider
        from agents.tracing import setup as tracing_setup

        tracing_setup._DEFAULT_SHUTDOWN_TIMEOUT = 0.2

        exporter = BackendSpanExporter(
            api_key="test_key",
            endpoint="http://127.0.0.1:{server.server_port}/traces/ingest",
            max_retries=100,
            base_delay=10.0,
            max_delay=10.0,
        )
        processor = BatchTraceProcessor(
            exporter=exporter,
            max_queue_size=1,
            max_batch_size=1,
            schedule_delay=60.0,
            export_trigger_ratio=1.0,
        )
        provider = DefaultTraceProvider()
        provider.register_processor(processor)
        original_shutdown = provider.shutdown

        def timed_shutdown(*args, **kwargs):
            shutdown_started = time.monotonic()
            try:
                return original_shutdown(*args, **kwargs)
            finally:
                print(
                    f"shutdown_elapsed={{time.monotonic() - shutdown_started:.6f}}",
                    flush=True,
                )

        provider.shutdown = timed_shutdown
        tracing_setup.set_trace_provider(provider)

        with trace("probe"):
            with custom_span("probe-span"):
                pass

        time.sleep(0.3)
        sys.exit(7)
        """
    )

    try:
        result = subprocess.run(
            [sys.executable, "-c", script],
            check=False,
            capture_output=True,
            text=True,
            timeout=10.0,
        )
    finally:
        server.shutdown()
        server.server_close()

    assert request_seen.is_set()
    assert result.returncode == 7
    shutdown_elapsed_prefix = "shutdown_elapsed="
    shutdown_elapsed_lines = [
        line for line in result.stdout.splitlines() if line.startswith(shutdown_elapsed_prefix)
    ]
    assert len(shutdown_elapsed_lines) == 1
    assert float(shutdown_elapsed_lines[0][len(shutdown_elapsed_prefix) :]) < 0.5


@patch("httpx.Client")
def test_backend_span_exporter_request_error(mock_client):
    # Make post() raise a RequestError each time
    mock_client.return_value.post.side_effect = httpx.RequestError("Network error")

    exporter = BackendSpanExporter(api_key="test_key", max_retries=2, base_delay=0.1, max_delay=0.2)
    with patch.object(exporter._shutdown_event, "wait", return_value=False) as wait_for_retry:
        exporter.export([get_span(mock_processor())])

    # Should retry up to max_retries times
    assert mock_client.return_value.post.call_count == 2
    wait_for_retry.assert_called_once()

    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_close(mock_client):
    exporter = BackendSpanExporter(api_key="test_key")
    exporter.close()

    # Ensure underlying http client is closed
    mock_client.return_value.close.assert_called_once()


@patch("httpx.Client")
def test_backend_span_exporter_sanitizes_generation_usage_for_openai_tracing(mock_client):
    """Unsupported usage keys should be stripped before POSTing to OpenAI tracing."""

    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            self.exported_payload: dict[str, Any] = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "usage": {
                        "requests": 1,
                        "input_tokens": 10,
                        "output_tokens": 5,
                        "total_tokens": 15,
                        "input_tokens_details": {"cached_tokens": 1},
                        "output_tokens_details": {"reasoning_tokens": 2},
                    },
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    item = DummyItem()
    exporter.export([cast(Any, item)])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    sent_usage = sent_payload["span_data"]["usage"]
    assert "requests" not in sent_usage
    assert "total_tokens" not in sent_usage
    assert "input_tokens_details" not in sent_usage
    assert "output_tokens_details" not in sent_usage
    assert sent_usage["input_tokens"] == 10
    assert sent_usage["output_tokens"] == 5
    assert sent_usage["details"] == {
        "requests": 1,
        "total_tokens": 15,
        "input_tokens_details": {"cached_tokens": 1},
        "output_tokens_details": {"reasoning_tokens": 2},
    }

    # Ensure the original exported object has not been mutated.
    assert "requests" in item.exported_payload["span_data"]["usage"]
    assert item.exported_payload["span_data"]["usage"]["total_tokens"] == 15
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_truncates_large_input_for_openai_tracing(mock_client):
    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            self.exported_payload: dict[str, Any] = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "input": "x" * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000),
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    item = DummyItem()
    exporter.export([cast(Any, item)])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    sent_input = sent_payload["span_data"]["input"]
    assert isinstance(sent_input, str)
    assert sent_input.endswith(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX)
    assert exporter._value_json_size_bytes(sent_input) <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    assert item.exported_payload["span_data"]["input"] != sent_input
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_truncates_large_structured_input_without_stringifying(mock_client):
    class NoStringifyDict(dict[str, Any]):
        def __str__(self) -> str:
            raise AssertionError("__str__ should not be called for oversized non-string previews")

    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            payload_input = NoStringifyDict(
                blob="x" * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000)
            )
            self.exported_payload: dict[str, Any] = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "input": payload_input,
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([cast(Any, DummyItem())])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    sent_input = sent_payload["span_data"]["input"]
    assert isinstance(sent_input, dict)
    assert isinstance(sent_input["blob"], str)
    assert sent_input["blob"].endswith(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX)
    assert exporter._value_json_size_bytes(sent_input) <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_keeps_generation_usage_for_custom_endpoint(mock_client):
    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            self.exported_payload = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "usage": {
                        "requests": 1,
                        "input_tokens": 10,
                        "output_tokens": 5,
                    },
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(
        api_key="test_key",
        endpoint="https://example.com/v1/traces/ingest",
    )
    exporter.export([cast(Any, DummyItem())])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    assert sent_payload["span_data"]["usage"]["requests"] == 1
    assert sent_payload["span_data"]["usage"]["input_tokens"] == 10
    assert sent_payload["span_data"]["usage"]["output_tokens"] == 5
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_drops_non_generation_usage_for_openai_endpoint(mock_client):
    class DummyItem:
        tracing_api_key = None

        def export(self):
            return {
                "object": "trace.span",
                "span_data": {
                    "type": "function",
                    "usage": {"requests": 1},
                },
            }

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(api_key="test_key")
    exporter.export([cast(Any, DummyItem())])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    assert "usage" not in sent_payload["span_data"]
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_keeps_non_generation_usage_for_custom_endpoint(mock_client):
    class DummyItem:
        tracing_api_key = None

        def export(self):
            return {
                "object": "trace.span",
                "span_data": {
                    "type": "function",
                    "usage": {"requests": 1},
                },
            }

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(
        api_key="test_key",
        endpoint="https://example.com/v1/traces/ingest",
    )
    exporter.export([cast(Any, DummyItem())])

    sent_payload = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    assert sent_payload["span_data"]["usage"] == {"requests": 1}
    exporter.close()


def test_sanitize_for_openai_tracing_api_keeps_allowed_generation_usage():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
            },
        },
    }
    assert exporter._sanitize_for_openai_tracing_api(payload) is payload
    exporter.close()


@patch("httpx.Client")
def test_backend_span_exporter_keeps_large_input_for_custom_endpoint(mock_client):
    class DummyItem:
        tracing_api_key = None

        def __init__(self):
            self.exported_payload: dict[str, Any] = {
                "object": "trace.span",
                "span_data": {
                    "type": "generation",
                    "input": "x" * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000),
                },
            }

        def export(self):
            return self.exported_payload

    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_client.return_value.post.return_value = mock_response

    exporter = BackendSpanExporter(
        api_key="test_key",
        endpoint="https://example.com/v1/traces/ingest",
    )
    item = DummyItem()
    exporter.export([cast(Any, item)])

    sent_payload: dict[str, Any] = mock_client.return_value.post.call_args.kwargs["json"]["data"][0]
    assert sent_payload["span_data"]["input"] == item.exported_payload["span_data"]["input"]
    exporter.close()


def test_sanitize_for_openai_tracing_api_moves_unsupported_generation_usage_to_details():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
                "total_tokens": 3,
                "input_tokens_details": {"cached_tokens": 0},
                "output_tokens_details": {"reasoning_tokens": 0},
                "details": {"provider": "litellm"},
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["usage"] == {
        "input_tokens": 1,
        "output_tokens": 2,
        "details": {
            "provider": "litellm",
            "total_tokens": 3,
            "input_tokens_details": {"cached_tokens": 0},
            "output_tokens_details": {"reasoning_tokens": 0},
        },
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_filters_non_json_values_in_usage_details():
    exporter = BackendSpanExporter(api_key="test_key")
    non_json = object()
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
                "input_tokens_details": {
                    "cached_tokens": 0,
                    "bad": non_json,
                },
                "output_tokens_details": {"reasoning_tokens": 0},
                "provider_usage": [1, non_json, {"ok": True, "bad": non_json}],
                "details": {
                    "provider": "litellm",
                    "bad": non_json,
                    "nested": {"keep": 1, "bad": non_json},
                },
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["usage"] == {
        "input_tokens": 1,
        "output_tokens": 2,
        "details": {
            "provider": "litellm",
            "nested": {"keep": 1},
            "input_tokens_details": {"cached_tokens": 0},
            "output_tokens_details": {"reasoning_tokens": 0},
            "provider_usage": [1, {"ok": True}],
        },
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_handles_cyclic_usage_values():
    exporter = BackendSpanExporter(api_key="test_key")
    cyclic_dict: dict[str, Any] = {}
    cyclic_dict["self"] = cyclic_dict
    cyclic_list: list[Any] = []
    cyclic_list.append(cyclic_list)

    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
                "input_tokens_details": cyclic_dict,
                "details": {
                    "provider": "litellm",
                    "cycle": cyclic_list,
                },
            },
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["usage"] == {
        "input_tokens": 1,
        "output_tokens": 2,
        "details": {
            "provider": "litellm",
            "cycle": [],
            "input_tokens_details": {},
        },
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_drops_non_dict_generation_usage_details():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "output_tokens": 2,
                "details": "invalid",
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["usage"] == {
        "input_tokens": 1,
        "output_tokens": 2,
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_drops_generation_usage_missing_required_tokens():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": 1,
                "total_tokens": 3,
                "input_tokens_details": {"cached_tokens": 0},
                "output_tokens_details": {"reasoning_tokens": 0},
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"] == {
        "type": "generation",
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_rejects_boolean_token_counts():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": {
                "input_tokens": True,
                "output_tokens": False,
                "input_tokens_details": {"cached_tokens": 0},
                "output_tokens_details": {"reasoning_tokens": 0},
            },
        },
    }
    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"] == {
        "type": "generation",
    }
    exporter.close()


def test_sanitize_for_openai_tracing_api_skips_non_dict_generation_usage():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "usage": None,
        },
    }
    assert exporter._sanitize_for_openai_tracing_api(payload) is payload
    exporter.close()


def test_sanitize_for_openai_tracing_api_keeps_small_input_without_mutation():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "input": "short input",
            "usage": {"input_tokens": 1, "output_tokens": 2},
        },
    }

    assert exporter._sanitize_for_openai_tracing_api(payload) is payload
    exporter.close()


def test_sanitize_for_openai_tracing_api_truncates_oversized_output():
    exporter = BackendSpanExporter(api_key="test_key")
    payload: dict[str, Any] = {
        "object": "trace.span",
        "span_data": {
            "type": "function",
            "output": "x" * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000),
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized is not payload
    assert sanitized["span_data"]["output"].endswith(
        exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX
    )
    assert (
        exporter._value_json_size_bytes(sanitized["span_data"]["output"])
        <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    )
    assert payload["span_data"]["output"] != sanitized["span_data"]["output"]
    exporter.close()


def test_sanitize_for_openai_tracing_api_preserves_generation_input_list_shape():
    exporter = BackendSpanExporter(api_key="test_key")
    payload = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "input": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "input_audio",
                            "input_audio": {
                                "data": "x"
                                * (BackendSpanExporter._OPENAI_TRACING_MAX_FIELD_BYTES + 5_000),
                                "format": "wav",
                            },
                        }
                    ],
                }
            ],
            "usage": {"input_tokens": 1, "output_tokens": 1},
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    sanitized_input = sanitized["span_data"]["input"]
    assert isinstance(sanitized_input, list)
    assert isinstance(sanitized_input[0], dict)
    assert sanitized_input[0]["role"] == "user"
    assert (
        exporter._value_json_size_bytes(sanitized_input) <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    )
    exporter.close()


def test_sanitize_for_openai_tracing_api_replaces_unserializable_output():
    exporter = BackendSpanExporter(api_key="test_key")
    payload: dict[str, Any] = {
        "object": "trace.span",
        "span_data": {
            "type": "function",
            "output": b"x" * 10,
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    assert sanitized["span_data"]["output"] == {
        "truncated": True,
        "original_type": "bytes",
        "preview": "<bytes bytes=10 truncated>",
    }
    exporter.close()


def test_truncate_json_value_for_limit_terminates_preview_dict_under_zero_budget():
    exporter = BackendSpanExporter(api_key="test_key")
    preview = exporter._truncated_preview(None)

    truncated = exporter._truncate_json_value_for_limit(preview, 0)

    assert truncated == {}
    exporter.close()


def test_sanitize_for_openai_tracing_api_handles_none_content_under_tight_budget():
    exporter = BackendSpanExporter(api_key="test_key")
    payload: dict[str, Any] = {
        "object": "trace.span",
        "span_data": {
            "type": "generation",
            "output": [
                {
                    "role": "assistant",
                    "content": None,
                    "name": "a" * 25_000,
                    "tool_calls": [],
                }
                for _ in range(8)
            ],
            "usage": {"input_tokens": 1, "output_tokens": 1},
        },
    }

    sanitized = exporter._sanitize_for_openai_tracing_api(payload)
    sanitized_output = cast(list[Any], sanitized["span_data"]["output"])

    assert isinstance(sanitized_output, list)
    assert sanitized_output != payload["span_data"]["output"]
    assert (
        exporter._value_json_size_bytes(sanitized_output)
        <= exporter._OPENAI_TRACING_MAX_FIELD_BYTES
    )
    assert any(item == {} for item in sanitized_output)
    exporter.close()


def test_truncate_string_for_json_limit_returns_original_when_within_limit():
    exporter = BackendSpanExporter(api_key="test_key")
    value = "hello"
    max_bytes = exporter._value_json_size_bytes(value)

    assert exporter._truncate_string_for_json_limit(value, max_bytes) == value
    exporter.close()


def test_truncate_string_for_json_limit_returns_suffix_when_limit_equals_suffix():
    exporter = BackendSpanExporter(api_key="test_key")
    max_bytes = exporter._value_json_size_bytes(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX)

    assert (
        exporter._truncate_string_for_json_limit("x" * 100, max_bytes)
        == exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX
    )
    exporter.close()


def test_truncate_string_for_json_limit_returns_empty_when_suffix_too_large():
    exporter = BackendSpanExporter(api_key="test_key")
    max_bytes = (
        exporter._value_json_size_bytes(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX) - 1
    )

    assert exporter._truncate_string_for_json_limit("x" * 100, max_bytes) == ""
    exporter.close()


def test_truncate_string_for_json_limit_handles_escape_heavy_input():
    exporter = BackendSpanExporter(api_key="test_key")
    value = ('\\"' * 40_000) + "tail"
    max_bytes = exporter._OPENAI_TRACING_MAX_FIELD_BYTES

    truncated = exporter._truncate_string_for_json_limit(value, max_bytes)

    assert truncated.endswith(exporter._OPENAI_TRACING_STRING_TRUNCATION_SUFFIX)
    assert exporter._value_json_size_bytes(truncated) <= max_bytes
    exporter.close()
