"""Tests for staso.provider (Span, Tracer, TracerProvider)."""
from __future__ import annotations

import json

from staso.provider import Span, Tracer, TracerProvider, _hex_to_uuid
from staso.types import SpanKind, SpanStatus


class TestHexToUuid:
    def test_full_hex(self):
        result = _hex_to_uuid("0123456789abcdef0123456789abcdef")
        assert result == "01234567-89ab-cdef-0123-456789abcdef"

    def test_short_hex_padded(self):
        result = _hex_to_uuid("abc")
        assert result.startswith("00000000-0000-0000-0000-000000000abc")

    def test_empty_string(self):
        result = _hex_to_uuid("")
        assert result == "00000000-0000-0000-0000-000000000000"


class TestSpan:
    def _make_span(self, transport=None, name="test-span", kind=SpanKind.CHAIN):
        return Span(
            name=name,
            trace_id="aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee",
            span_id="11111111-2222-3333-4444-555555555555",
            parent_span_id=None,
            transport=transport,
            agent_id="test-agent",
            session_id="test-session",
            kind=kind,
        )

    def test_recording_flag(self):
        span = self._make_span()
        assert span._recording is True

    def test_not_recording_after_end(self):
        span = self._make_span()
        span.end()
        assert span._recording is False

    def test_set_status(self):
        span = self._make_span()
        span.set_status(SpanStatus.OK)
        assert span.status == SpanStatus.OK

    def test_set_status_error_with_description(self):
        span = self._make_span()
        span.set_status(SpanStatus.ERROR, "something broke")
        assert span.status == SpanStatus.ERROR
        assert span.error_message == "something broke"

    def test_set_status_ignored_after_end(self):
        span = self._make_span()
        span.end()
        span.set_status(SpanStatus.ERROR, "late")
        assert span.status == SpanStatus.OK  # default, unchanged

    def test_record_exception(self):
        span = self._make_span()
        span.record_exception(ValueError("bad value"))
        assert span.status == SpanStatus.ERROR
        assert span.error_message == "bad value"

    def test_end_enqueues_to_transport(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(transport=transport, kind=SpanKind.AGENT)
        span.end()

        transport.enqueue.assert_called_once()
        span_dict = transport.enqueue.call_args[0][0]
        assert span_dict["name"] == "test-span"
        assert span_dict["kind"] == "agent"
        assert span_dict["status"] == "ok"
        assert span_dict["session_id"] == "test-session"
        assert span_dict["agent_id"] == "test-agent"
        assert span_dict["error_message"] is None

    def test_end_captures_error_message(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(transport=transport)
        span.set_status(SpanStatus.ERROR, "failed")
        span.end()

        span_dict = transport.enqueue.call_args[0][0]
        assert span_dict["status"] == "error"
        assert span_dict["error_message"] == "failed"

    def test_end_captures_llm_fields(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(transport=transport, kind=SpanKind.LLM)
        span.model = "claude-sonnet-4-20250514"
        span.input_tokens = 100
        span.output_tokens = 50
        span.total_tokens = 150
        span.input = {"prompt": "hello"}
        span.output = {"response": "hi"}
        span.end()

        d = transport.enqueue.call_args[0][0]
        assert d["kind"] == "llm"
        assert d["model"] == "claude-sonnet-4-20250514"
        assert d["input_tokens"] == 100
        assert d["output_tokens"] == 50
        assert d["total_tokens"] == 150
        assert json.loads(d["input"]) == {"prompt": "hello"}
        assert json.loads(d["output"]) == {"response": "hi"}

    def test_end_puts_metadata_in_attributes(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(transport=transport)
        span.metadata["tags"] = ["a", "b"]
        span.end()

        d = transport.enqueue.call_args[0][0]
        attrs = json.loads(d["attributes"])
        assert attrs["tags"] == ["a", "b"]

    def test_end_no_double_enqueue(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(transport=transport)
        span.end()
        span.end()  # second end should be no-op
        assert transport.enqueue.call_count == 1

    def test_end_with_no_transport(self):
        span = self._make_span(transport=None)
        span.end()  # no error

    def test_duration_ms_positive(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(transport=transport)
        span.end()

        d = transport.enqueue.call_args[0][0]
        assert d["duration_ms"] >= 0

    def test_timestamp_is_iso8601(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(transport=transport)
        span.end()

        d = transport.enqueue.call_args[0][0]
        assert "T" in d["timestamp"]


class TestTracer:
    def _make_provider(self, transport=None):
        return TracerProvider(transport=transport, agent_id="test-agent")

    def test_start_span_returns_span(self):
        provider = self._make_provider()
        tracer = provider.get_tracer()
        span = tracer.start_span("my-span")
        assert isinstance(span, Span)
        span.end()

    def test_start_span_generates_ids(self):
        provider = self._make_provider()
        tracer = provider.get_tracer()
        span = tracer.start_span("my-span")
        assert span.trace_id != ""
        assert span.span_id != ""
        span.end()

    def test_child_span_inherits_trace_id(self):
        provider = self._make_provider()
        tracer = provider.get_tracer()

        with tracer.start_as_current_span("parent") as parent:
            child = tracer.start_span("child")
            assert child.trace_id == parent.trace_id
            child.end()

    def test_child_span_has_parent_id(self):
        provider = self._make_provider()
        tracer = provider.get_tracer()

        with tracer.start_as_current_span("parent") as parent:
            child = tracer.start_span("child")
            assert child.parent_span_id == parent.span_id
            child.end()

    def test_start_as_current_span_records_exception(self, mocker):
        transport = mocker.MagicMock()
        provider = self._make_provider(transport=transport)
        tracer = provider.get_tracer()

        import pytest

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

        transport.enqueue.assert_called_once()
        d = transport.enqueue.call_args[0][0]
        assert d["status"] == "error"
        assert d["error_message"] == "boom"


class TestTracerProvider:
    def test_get_tracer_returns_tracer(self):
        provider = TracerProvider(transport=None, agent_id="test")
        tracer = provider.get_tracer()
        assert isinstance(tracer, Tracer)

    def test_shutdown_calls_transport_shutdown(self, mocker):
        transport = mocker.MagicMock()
        provider = TracerProvider(transport=transport, agent_id="test")
        provider.shutdown()
        transport.shutdown.assert_called_once()

    def test_shutdown_no_transport(self):
        provider = TracerProvider(transport=None, agent_id="test")
        provider.shutdown()  # should not raise
