import asyncio

from typing import NoReturn
from unittest import mock

import pytest

from a2a.utils.telemetry import trace_class, trace_function


@pytest.fixture
def mock_span():
    return mock.MagicMock()


@pytest.fixture
def mock_tracer(mock_span):
    tracer = mock.MagicMock()
    tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
    tracer.start_as_current_span.return_value.__exit__.return_value = False
    return tracer


@pytest.fixture(autouse=True)
def patch_trace_get_tracer(mock_tracer):
    with mock.patch('opentelemetry.trace.get_tracer', return_value=mock_tracer):
        yield


def test_trace_function_sync_success(mock_span):
    @trace_function
    def foo(x, y):
        return x + y

    result = foo(2, 3)
    assert result == 5
    mock_span.set_status.assert_called()
    mock_span.set_status.assert_any_call(mock.ANY)
    mock_span.record_exception.assert_not_called()


def test_trace_function_sync_exception(mock_span):
    @trace_function
    def bar() -> NoReturn:
        raise ValueError('fail')

    with pytest.raises(ValueError):
        bar()
    mock_span.record_exception.assert_called()
    mock_span.set_status.assert_any_call(mock.ANY, description='fail')


def test_trace_function_sync_attribute_extractor_called(mock_span):
    called = {}

    def attr_extractor(span, args, kwargs, result, exception) -> None:
        called['called'] = True
        assert span is mock_span
        assert exception is None
        assert result == 42

    @trace_function(attribute_extractor=attr_extractor)
    def foo() -> int:
        return 42

    foo()
    assert called['called']


def test_trace_function_sync_attribute_extractor_error_logged(mock_span):
    with mock.patch('a2a.utils.telemetry.logger') as logger:

        def attr_extractor(span, args, kwargs, result, exception) -> NoReturn:
            raise RuntimeError('attr fail')

        @trace_function(attribute_extractor=attr_extractor)
        def foo() -> int:
            return 1

        foo()
        logger.error.assert_any_call(mock.ANY)


@pytest.mark.asyncio
async def test_trace_function_async_success(mock_span):
    @trace_function
    async def foo(x):
        await asyncio.sleep(0)
        return x * 2

    result = await foo(4)
    assert result == 8
    mock_span.set_status.assert_called()
    mock_span.record_exception.assert_not_called()


@pytest.mark.asyncio
async def test_trace_function_async_exception(mock_span):
    @trace_function
    async def bar() -> NoReturn:
        await asyncio.sleep(0)
        raise RuntimeError('async fail')

    with pytest.raises(RuntimeError):
        await bar()
    mock_span.record_exception.assert_called()
    mock_span.set_status.assert_any_call(mock.ANY, description='async fail')


@pytest.mark.asyncio
async def test_trace_function_async_attribute_extractor_called(mock_span):
    called = {}

    def attr_extractor(span, args, kwargs, result, exception) -> None:
        called['called'] = True
        assert exception is None
        assert result == 99

    @trace_function(attribute_extractor=attr_extractor)
    async def foo() -> int:
        return 99

    await foo()
    assert called['called']


def test_trace_function_with_args_and_attributes(mock_span):
    @trace_function(span_name='custom.span', attributes={'foo': 'bar'})
    def foo() -> int:
        return 1

    foo()
    mock_span.set_attribute.assert_any_call('foo', 'bar')


def test_trace_class_exclude_list(mock_span):
    @trace_class(exclude_list=['skip_me'])
    class MyClass:
        def a(self) -> str:
            return 'a'

        def skip_me(self) -> str:
            return 'skip'

        def __str__(self):
            return 'str'

    obj = MyClass()
    assert obj.a() == 'a'
    assert obj.skip_me() == 'skip'
    # Only 'a' is traced, not 'skip_me' or dunder
    assert hasattr(obj.a, '__wrapped__')
    assert not hasattr(obj.skip_me, '__wrapped__')


def test_trace_class_include_list(mock_span):
    @trace_class(include_list=['only_this'])
    class MyClass:
        def only_this(self) -> str:
            return 'yes'

        def not_this(self) -> str:
            return 'no'

    obj = MyClass()
    assert obj.only_this() == 'yes'
    assert obj.not_this() == 'no'
    assert hasattr(obj.only_this, '__wrapped__')
    assert not hasattr(obj.not_this, '__wrapped__')


def test_trace_class_dunder_not_traced(mock_span):
    @trace_class()
    class MyClass:
        def __init__(self):
            self.x = 1

        def foo(self) -> str:
            return 'foo'

    obj = MyClass()
    assert obj.foo() == 'foo'
    assert hasattr(obj.foo, '__wrapped__')
    assert hasattr(obj, 'x')