"""Tests for st.annotate() and span.annotate()."""
from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

import staso as st
from staso.client import _set_client
from staso.annotation import AnnotationDataType, shutdown_annotation_client


class TestAnnotate:
    def test_annotate_sends_http_post(self, capturing, mocker):
        mock_client = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = 201
        mock_client.post.return_value = mock_response

        with patch("staso.annotation._get_annotation_client", return_value=mock_client):
            st.annotate(trace_id="abc", name="accuracy", value=0.95)

        mock_client.post.assert_called_once()
        call_args = mock_client.post.call_args
        assert "/v1/traces/abc/annotations" in call_args[0][0]
        payload = call_args[1]["json"]
        assert payload["name"] == "accuracy"
        assert payload["value"] == 0.95
        assert payload["data_type"] == "numeric"

    def test_annotate_with_all_fields(self, capturing, mocker):
        mock_client = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = 201
        mock_client.post.return_value = mock_response

        with patch("staso.annotation._get_annotation_client", return_value=mock_client):
            st.annotate(
                trace_id="abc",
                name="quality",
                value=0.8,
                data_type=AnnotationDataType.NUMERIC,
                span_id="span-1",
                comment="Good response",
            )

        payload = mock_client.post.call_args[1]["json"]
        assert payload["name"] == "quality"
        assert payload["value"] == 0.8
        assert payload["data_type"] == "numeric"
        assert payload["span_id"] == "span-1"
        assert payload["comment"] == "Good response"

    def test_annotate_boolean_type(self, capturing, mocker):
        mock_client = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = 201
        mock_client.post.return_value = mock_response

        with patch("staso.annotation._get_annotation_client", return_value=mock_client):
            st.annotate(
                trace_id="abc", name="correct", value=True, data_type="boolean"
            )

        payload = mock_client.post.call_args[1]["json"]
        assert payload["value"] is True
        assert payload["data_type"] == "boolean"

    def test_annotate_categorical_type(self, capturing, mocker):
        mock_client = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = 201
        mock_client.post.return_value = mock_response

        with patch("staso.annotation._get_annotation_client", return_value=mock_client):
            st.annotate(
                trace_id="abc",
                name="sentiment",
                value="positive",
                data_type="categorical",
            )

        payload = mock_client.post.call_args[1]["json"]
        assert payload["value"] == "positive"
        assert payload["data_type"] == "categorical"

    def test_annotate_noop_without_init(self, mocker):
        _set_client(None)  # type: ignore[arg-type]
        mock_post = mocker.patch("httpx.Client.post")

        st.annotate(trace_id="abc", name="x", value=1.0)

        mock_post.assert_not_called()

    def test_annotate_validation_empty_name(self, capturing):
        with pytest.raises(ValueError, match="non-empty string"):
            st.annotate(trace_id="abc", name="", value=1.0)

    def test_span_annotate_convenience(self, capturing, mocker):
        mock_client = MagicMock()
        mock_response = MagicMock()
        mock_response.status_code = 201
        mock_client.post.return_value = mock_response

        with patch("staso.annotation._get_annotation_client", return_value=mock_client):
            tracer = st.get_client().provider.get_tracer()
            with tracer.start_as_current_span("test") as span:
                span.annotate("quality", value=0.8)

        mock_client.post.assert_called_once()
        call_args = mock_client.post.call_args
        assert f"/v1/traces/{span.trace_id}/annotations" in call_args[0][0]
        payload = call_args[1]["json"]
        assert payload["name"] == "quality"
        assert payload["value"] == 0.8
        assert payload["span_id"] == span.span_id
