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

import json

from openinference.semconv.trace import SpanAttributes as OI
from opentelemetry.trace import SpanContext, StatusCode, TraceFlags

from staso.provider import Span, Tracer, TracerProvider, _hex_to_uuid, _OI_KIND_MAP


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 TestOiKindMap:
    def test_all_kinds_mapped(self):
        assert _OI_KIND_MAP["AGENT"] == "agent"
        assert _OI_KIND_MAP["LLM"] == "llm"
        assert _OI_KIND_MAP["TOOL"] == "tool"
        assert _OI_KIND_MAP["CHAIN"] == "chain"
        assert _OI_KIND_MAP["EMBEDDING"] == "llm"
        assert _OI_KIND_MAP["RETRIEVER"] == "tool"


class TestSpan:
    def _make_span(self, transport=None, name="test-span", attributes=None):
        ctx = SpanContext(
            trace_id=1,
            span_id=2,
            is_remote=False,
            trace_flags=TraceFlags(TraceFlags.SAMPLED),
        )
        return Span(
            name=name,
            span_context=ctx,
            parent_span_id=None,
            transport=transport,
            agent_id="test-agent",
            session_id="test-session",
            attributes=attributes,
        )

    def test_is_recording(self):
        span = self._make_span()
        assert span.is_recording() is True

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

    def test_get_span_context(self):
        span = self._make_span()
        ctx = span.get_span_context()
        assert ctx.trace_id == 1
        assert ctx.span_id == 2

    def test_set_attribute(self):
        span = self._make_span()
        span.set_attribute("key", "value")
        assert span._attributes["key"] == "value"

    def test_set_attribute_ignored_after_end(self):
        span = self._make_span()
        span.end()
        span.set_attribute("key", "value")
        assert "key" not in span._attributes

    def test_set_attributes_batch(self):
        span = self._make_span()
        span.set_attributes({"a": 1, "b": 2})
        assert span._attributes["a"] == 1
        assert span._attributes["b"] == 2

    def test_update_name(self):
        span = self._make_span(name="old")
        span.update_name("new")
        assert span._name == "new"

    def test_update_name_ignored_after_end(self):
        span = self._make_span(name="old")
        span.end()
        span.update_name("new")
        assert span._name == "old"

    def test_set_status(self):
        span = self._make_span()
        span.set_status(StatusCode.OK)
        assert span._status_code == StatusCode.OK

    def test_set_status_error_with_description(self):
        span = self._make_span()
        span.set_status(StatusCode.ERROR, "something broke")
        assert span._status_code == StatusCode.ERROR
        assert span._status_description == "something broke"

    def test_record_exception(self):
        span = self._make_span()
        span.record_exception(ValueError("bad value"))
        assert span._status_code == StatusCode.ERROR
        assert span._attributes["exception.type"] == "ValueError"
        assert span._attributes["exception.message"] == "bad value"

    def test_end_enqueues_to_transport(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(
            transport=transport,
            attributes={OI.OPENINFERENCE_SPAN_KIND: "AGENT"},
        )
        span.set_status(StatusCode.OK)
        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(StatusCode.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_extracts_llm_fields(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(
            transport=transport,
            attributes={
                OI.OPENINFERENCE_SPAN_KIND: "LLM",
                OI.LLM_MODEL_NAME: "claude-sonnet-4-20250514",
                OI.LLM_PROVIDER: "anthropic",
                OI.LLM_TOKEN_COUNT_PROMPT: 100,
                OI.LLM_TOKEN_COUNT_COMPLETION: 50,
                OI.LLM_TOKEN_COUNT_TOTAL: 150,
                OI.INPUT_VALUE: '{"prompt": "hello"}',
                OI.OUTPUT_VALUE: '{"response": "hi"}',
            },
        )
        span.set_status(StatusCode.OK)
        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 d["input"] == '{"prompt": "hello"}'
        assert d["output"] == '{"response": "hi"}'

    def test_end_puts_tags_in_attributes(self, mocker):
        transport = mocker.MagicMock()
        span = self._make_span(
            transport=transport,
            attributes={OI.OPENINFERENCE_SPAN_KIND: "TOOL", OI.TAG_TAGS: ["a", "b"]},
        )
        span.set_status(StatusCode.OK)
        span.end()

        d = transport.enqueue.call_args[0][0]
        remaining = json.loads(d["attributes"])
        assert remaining["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 with no transport (disabled SDK) should not raise."""
        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.set_status(StatusCode.OK)
        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]
        # ISO 8601 contains 'T' separator
        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("test")
        span = tracer.start_span("my-span")
        assert isinstance(span, Span)
        span.end()

    def test_start_span_generates_valid_context(self):
        provider = self._make_provider()
        tracer = provider.get_tracer("test")
        span = tracer.start_span("my-span")
        ctx = span.get_span_context()
        assert ctx.trace_id != 0
        assert ctx.span_id != 0
        span.end()

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

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

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

        with tracer.start_as_current_span("parent") as parent:
            child = tracer.start_span("child")
            assert isinstance(child, Span)
            assert child._parent_span_id == parent.get_span_context().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("test")

        import pytest

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

        # Span should have been ended and enqueued
        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("module-name")
        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
