from unittest import mock

import pytest
from openinference.semconv.trace import SpanAttributes

from cancontech_tracing import traceable


@pytest.fixture
def mock_tracer():
    with mock.patch("opentelemetry.trace.get_tracer") as mock_get_tracer:
        mock_tracer = mock.MagicMock()
        mock_get_tracer.return_value = mock_tracer

        mock_span = mock.Mock()
        mock_set_current_span = mock.MagicMock()
        mock_set_current_span.__enter__ = mock.Mock(return_value=mock_span)

        mock_use_span = mock.MagicMock()
        mock_use_span.__enter__ = mock.Mock(return_value=mock_span)

        mock_tracer.start_span = mock.Mock(return_value=mock_span)
        mock_tracer.use_span = mock.Mock(return_value=mock_use_span)

        yield mock_tracer, mock_span


@traceable()
def tracing_target_func(x, y):
    return x + y


@traceable(
    name="test_name",
    get_input=lambda x, *args: f"x: {x}",
    get_output=lambda res: "output_test",
)
def tracing_target_func_with_kwargs_input(x=5, y=6):
    return x + y


class CustomInputObject:
    def __init__(self, name):
        self.name = name


@traceable()
def tracing_target_custom_object(data: CustomInputObject):
    return data.name


@traceable()
def tracing_target_generator():
    for i in range(3):
        yield i


@traceable()
async def tracing_target_async_generator():
    for i in range(3):
        yield i


def test_traceable_with_default_properties(mock_tracer):
    mock_tracer_instance, mock_span = mock_tracer
    x = 10
    y = 15

    result = tracing_target_func(x, y)
    assert result == 25

    mock_tracer_instance.start_span.assert_called_once()
    span_args, span_kwargs = mock_tracer_instance.start_span.call_args

    assert span_kwargs["name"] == "tracing_target_func"
    assert span_kwargs["attributes"][SpanAttributes.INPUT_VALUE] == '{"x": 10, "y": 15}'

    # @TODO test OUTPUT_VALUE
    # mock_span.set_attribute.assert_called_with(SpanAttributes.OUTPUT_VALUE, "25")


def test_traceable_with_properties_set(mock_tracer):
    mock_tracer_instance, mock_span = mock_tracer
    x = 10
    y = 15

    result = tracing_target_func_with_kwargs_input(x, y)
    assert result == 25

    mock_tracer_instance.start_span.assert_called_once()
    span_args, span_kwargs = mock_tracer_instance.start_span.call_args

    assert span_kwargs["name"] == "test_name"
    assert span_kwargs["attributes"][SpanAttributes.INPUT_VALUE] == "x: 10"

    # @TODO test OUTPUT_VALUE
    # mock_span.set_attribute.assert_called_with(
    #     SpanAttributes.OUTPUT_VALUE, "output_test"
    # )


def test_traceable_with_custom_object_as_input(mock_tracer):
    mock_tracer_instance, mock_span = mock_tracer
    obj = CustomInputObject("test_name")

    result = tracing_target_custom_object(obj)
    assert result == "test_name"

    span_args, span_kwargs = mock_tracer_instance.start_span.call_args

    assert span_kwargs["name"] == "tracing_target_custom_object"
    assert "CustomInputObject object at" in span_kwargs["attributes"][SpanAttributes.INPUT_VALUE]


def test_traceable_with_generator(mock_tracer):
    mock_tracer_instance, mock_span = mock_tracer

    values = list(tracing_target_generator())
    assert values == [0, 1, 2]

    mock_tracer_instance.start_span.assert_called_once()
    span_args, span_kwargs = mock_tracer_instance.start_span.call_args

    assert span_kwargs["name"] == "tracing_target_generator"
    mock_span.add_event.assert_called_with("yielded_value")
    assert mock_span.add_event.call_count == 3


@pytest.mark.asyncio
async def test_traceable_with_async_generator(mock_tracer):
    mock_tracer_instance, mock_span = mock_tracer

    values = []
    async for value in tracing_target_async_generator():
        values.append(value)
    assert values == [0, 1, 2]

    mock_tracer_instance.start_span.assert_called_once()
    span_args, span_kwargs = mock_tracer_instance.start_span.call_args

    assert span_kwargs["name"] == "tracing_target_async_generator"
    mock_span.add_event.assert_called_with("yielded_value")
    assert mock_span.add_event.call_count == 3
