"""Tests for async function support in code validator."""

import pytest

from core.languages.python.testgen.postprocessing.code_validator import (
    CodeValidationError,
    split_code_with_ast,
    split_code_with_regex,
    validate_testgen_code,
)


def test_async_test_function_recognition() -> None:
    async_test_code = '''
import pytest
import asyncio

async def hello_world():
    print("Hello")
    await asyncio.sleep(0.1)
    return "World"

@pytest.mark.asyncio
async def test_hello_world_basic():
    """Basic async test function."""
    result = await hello_world()
    assert result == "World"
    '''

    result = validate_testgen_code(async_test_code, (3, 11))
    assert "async def test_hello_world_basic" in result
    assert "@pytest.mark.asyncio" in result


def test_mixed_sync_async_test_functions() -> None:
    mixed_test_code = '''
import pytest
import asyncio

def regular_function():
    return "sync"

async def async_function():
    await asyncio.sleep(0.1)
    return "async"

def test_sync_function():
    """Sync test function."""
    result = regular_function()
    assert result == "sync"

@pytest.mark.asyncio
async def test_async_function():
    """Async test function."""
    result = await async_function()
    assert result == "async"
    '''

    result = validate_testgen_code(mixed_test_code, (3, 11))
    assert "def test_sync_function" in result
    assert "async def test_async_function" in result


def test_async_unittest_style() -> None:
    async_unittest_code = '''
import unittest
import asyncio

class TestAsync(unittest.TestCase):
    async def test_async_method(self):
        """Async test method in unittest class."""
        await asyncio.sleep(0.1)
        self.assertEqual(1, 1)

    def test_sync_method(self):
        """Sync test method in unittest class."""
        self.assertEqual(2, 2)
    '''

    result = validate_testgen_code(async_unittest_code, (3, 11))
    assert "async def test_async_method" in result
    assert "def test_sync_method" in result


def test_no_async_test_functions_should_fail() -> None:
    no_test_code = '''
import asyncio

async def helper_function():
    """This is not a test function."""
    await asyncio.sleep(0.1)
    return "helper"

async def another_helper():
    """Another non-test async function."""
    return await helper_function()
    '''

    with pytest.raises(CodeValidationError, match="No test functions were found"):
        validate_testgen_code(no_test_code, (3, 11))


def test_async_function_with_complex_decorators() -> None:
    complex_decorator_code = '''
import pytest
import asyncio
from unittest.mock import patch

@pytest.mark.asyncio
@patch('some.module')
async def test_with_multiple_decorators(mock_obj):
    """Async test with multiple decorators."""
    await asyncio.sleep(0.1)
    assert True
    '''

    result = validate_testgen_code(complex_decorator_code, (3, 11))
    assert "async def test_with_multiple_decorators" in result
    assert "@pytest.mark.asyncio" in result
    assert "@patch" in result


class TestDecoratorHandling:
    """Tests for decorator handling in code splitting."""

    def test_split_code_with_ast_preserves_decorators(self) -> None:
        """Test that AST splitting preserves decorators with async test functions."""
        code = """import pytest

@pytest.mark.asyncio
async def test_foo():
    assert True

@pytest.mark.asyncio
@pytest.mark.parametrize("x", [1, 2])
async def test_bar(x):
    assert x > 0
"""
        result = split_code_with_ast(code, (3, 11))
        assert result is not None
        assert result.preamble == "import pytest\n"
        assert result.test_functions == [
            "@pytest.mark.asyncio\nasync def test_foo():\n    assert True\n",
            '@pytest.mark.asyncio\n@pytest.mark.parametrize("x", [1, 2])\nasync def test_bar(x):\n    assert x > 0\n',
        ]

    def test_split_code_with_regex_preserves_decorators(self) -> None:
        """Test that regex splitting preserves decorators with async test functions."""
        code = """import pytest

@pytest.mark.asyncio
async def test_foo():
    assert True

@pytest.mark.asyncio
@pytest.mark.parametrize("x", [1, 2])
async def test_bar(x):
    assert x > 0
"""
        result = split_code_with_regex(code)
        assert result.preamble == "import pytest\n"
        assert result.test_functions == [
            "@pytest.mark.asyncio\nasync def test_foo():\n    assert True",
            '@pytest.mark.asyncio\n@pytest.mark.parametrize("x", [1, 2])\nasync def test_bar(x):\n    assert x > 0',
        ]

    def test_validate_testgen_code_preserves_decorators(self) -> None:
        """End-to-end test that decorators are preserved through validation."""
        code = '''import pytest

@pytest.mark.asyncio
async def test_async_function():
    """Test async function."""
    assert True

def test_sync_function():
    """Test sync function."""
    assert True
'''
        result = validate_testgen_code(code, (3, 11))
        expected = '''import pytest


@pytest.mark.asyncio
async def test_async_function():
    """Test async function."""
    assert True


def test_sync_function():
    """Test sync function."""
    assert True'''
        assert result == expected
