"""Tests for staso.transport (BatchTransport, _spans_to_traces)."""
from __future__ import annotations

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

import pytest

from staso.transport import BatchTransport, _spans_to_traces


class TestSpansToTraces:
    def test_single_trace(self):
        spans = [
            {"trace_id": "t1", "span_id": "s1", "timestamp": "2026-01-01T00:00:01Z", "name": "a"},
            {"trace_id": "t1", "span_id": "s2", "timestamp": "2026-01-01T00:00:00Z", "name": "b"},
        ]
        traces = _spans_to_traces(spans, environment="prod", workspace_slug="ws")

        assert len(traces) == 1
        t = traces[0]
        assert t["trace_id"] == "t1"
        assert t["timestamp"] == "2026-01-01T00:00:00Z"  # earliest
        assert t["environment"] == "prod"
        assert t["workspace_slug"] == "ws"
        assert len(t["spans"]) == 2

    def test_multiple_traces(self):
        spans = [
            {"trace_id": "t1", "span_id": "s1", "timestamp": "2026-01-01T00:00:00Z"},
            {"trace_id": "t2", "span_id": "s2", "timestamp": "2026-01-01T00:00:01Z"},
        ]
        traces = _spans_to_traces(spans)
        assert len(traces) == 2
        trace_ids = {t["trace_id"] for t in traces}
        assert trace_ids == {"t1", "t2"}

    def test_empty_spans(self):
        assert _spans_to_traces([]) == []

    def test_span_fields_extracted(self):
        spans = [
            {
                "trace_id": "t1",
                "span_id": "s1",
                "parent_span_id": None,
                "name": "test-span",
                "kind": "agent",
                "timestamp": "2026-01-01T00:00:00Z",
                "duration_ms": 42.5,
                "status": "ok",
                "model": "claude",
                "input_tokens": 10,
                "output_tokens": 20,
                "total_tokens": 30,
                "input": '{"q": "hi"}',
                "output": '{"a": "hello"}',
                "attributes": "{}",
                "error_message": None,
                "session_id": "sess-1",
                "agent_id": "agent-1",
            }
        ]
        traces = _spans_to_traces(spans)
        s = traces[0]["spans"][0]

        # session_id and agent_id ARE included in the trace span
        assert s["session_id"] == "sess-1"
        assert s["agent_id"] == "agent-1"
        # And these should also be present
        assert s["name"] == "test-span"
        assert s["kind"] == "agent"
        assert s["duration_ms"] == 42.5
        assert s["model"] == "claude"
        assert s["input_tokens"] == 10

    def test_defaults_for_missing_fields(self):
        spans = [{"trace_id": "t1"}]
        traces = _spans_to_traces(spans)
        s = traces[0]["spans"][0]
        assert s["name"] == ""
        assert s["kind"] == "chain"
        assert s["status"] == "ok"
        assert s["model"] == ""
        assert s["input_tokens"] == 0


class TestBatchTransport:
    def test_constructs_url(self):
        with patch.object(BatchTransport, "_worker"):
            t = BatchTransport.__new__(BatchTransport)
            t._url = "http://localhost:6999/v1/traces/ingest"
        # Verify URL construction logic
        url = f"{'http://localhost:6999/'.rstrip('/')}/v1/traces/ingest"
        assert url == "http://localhost:6999/v1/traces/ingest"

    def test_enqueue_and_shutdown(self):
        """Integration: enqueue a span, shut down, verify flush was attempted."""
        with patch("staso.transport.httpx.Client") as mock_client_cls:
            mock_resp = MagicMock()
            mock_resp.status_code = 200
            mock_client_cls.return_value.post.return_value = mock_resp

            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
                batch_size=1,
                flush_interval=0.1,
            )

            transport.enqueue({
                "trace_id": "t1",
                "span_id": "s1",
                "timestamp": "2026-01-01T00:00:00Z",
                "name": "test",
            })

            # Give worker time to flush
            time.sleep(0.5)
            transport.shutdown()

            mock_client_cls.return_value.post.assert_called()
            call_kwargs = mock_client_cls.return_value.post.call_args
            assert "/v1/traces/ingest" in call_kwargs[0][0]

    def test_enqueue_after_shutdown_is_noop(self):
        with patch("staso.transport.httpx.Client"):
            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
                flush_interval=60,
            )
            transport.shutdown()
            # Should not raise
            transport.enqueue({"trace_id": "t1", "span_id": "s1"})

    def test_queue_full_drops_silently(self):
        with patch("staso.transport.httpx.Client"):
            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
                max_queue_size=1,
                flush_interval=60,
                batch_size=999,
            )
            # Fill the queue
            transport.enqueue({"trace_id": "t1"})
            # This should drop silently, not raise
            transport.enqueue({"trace_id": "t2"})
            transport.shutdown()

    def test_batch_size_capped_at_500(self):
        with patch("staso.transport.httpx.Client"):
            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
                batch_size=1000,
            )
            assert transport._batch_size == 500
            transport.shutdown()

    def test_flush_handles_http_error(self):
        """Flush should log but not raise on HTTP errors."""
        with patch("staso.transport.httpx.Client") as mock_client_cls:
            mock_resp = MagicMock()
            mock_resp.status_code = 500
            mock_resp.text = "Internal Server Error"
            mock_client_cls.return_value.post.return_value = mock_resp

            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
                batch_size=1,
                flush_interval=0.1,
            )
            transport.enqueue({
                "trace_id": "t1",
                "span_id": "s1",
                "timestamp": "2026-01-01T00:00:00Z",
            })
            time.sleep(0.5)
            transport.shutdown()  # Should not raise

    def test_flush_handles_network_exception(self):
        with patch("staso.transport.httpx.Client") as mock_client_cls:
            mock_client_cls.return_value.post.side_effect = ConnectionError("offline")

            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
                batch_size=1,
                flush_interval=0.1,
            )
            transport.enqueue({
                "trace_id": "t1",
                "span_id": "s1",
                "timestamp": "2026-01-01T00:00:00Z",
            })
            time.sleep(0.5)
            transport.shutdown()  # Should not raise

    def test_double_shutdown_is_safe(self):
        with patch("staso.transport.httpx.Client"):
            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
            )
            transport.shutdown()
            transport.shutdown()  # Second call should be no-op


class TestRetry:
    def test_retries_on_5xx(self):
        """Should retry up to 3 times on 5xx responses."""
        with patch("staso.transport.httpx.Client") as mock_client_cls:
            responses = [
                MagicMock(status_code=500, text="error"),
                MagicMock(status_code=503, text="unavailable"),
                MagicMock(status_code=200),
            ]
            mock_client_cls.return_value.post.side_effect = responses

            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
                batch_size=1,
                flush_interval=0.1,
            )
            transport.enqueue({
                "trace_id": "t1", "span_id": "s1",
                "timestamp": "2026-01-01T00:00:00Z",
            })
            time.sleep(1.0)
            transport.shutdown()

            assert mock_client_cls.return_value.post.call_count == 3

    def test_no_retry_on_4xx(self):
        """Should NOT retry on 4xx client errors."""
        with patch("staso.transport.httpx.Client") as mock_client_cls:
            mock_resp = MagicMock(status_code=400, text="bad request")
            mock_client_cls.return_value.post.return_value = mock_resp

            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
                batch_size=1,
                flush_interval=0.1,
            )
            transport.enqueue({
                "trace_id": "t1", "span_id": "s1",
                "timestamp": "2026-01-01T00:00:00Z",
            })
            time.sleep(0.5)
            transport.shutdown()

            # Only 1 call — no retry for 4xx
            assert mock_client_cls.return_value.post.call_count == 1

    def test_retries_on_network_error(self):
        """Should retry on network exceptions."""
        with patch("staso.transport.httpx.Client") as mock_client_cls:
            mock_client_cls.return_value.post.side_effect = [
                ConnectionError("offline"),
                ConnectionError("still offline"),
                MagicMock(status_code=200),
            ]

            transport = BatchTransport(
                base_url="http://localhost:6999",
                api_key="ak_test",
                batch_size=1,
                flush_interval=0.1,
            )
            transport.enqueue({
                "trace_id": "t1", "span_id": "s1",
                "timestamp": "2026-01-01T00:00:00Z",
            })
            time.sleep(2.0)
            transport.shutdown()

            assert mock_client_cls.return_value.post.call_count == 3
