from __future__ import annotations

import pytest

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


def test_valid_code() -> None:
    code: str = """def test_function():
    x = 1
    y = 2
    return x + y"""

    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result.strip() == code.strip()


def test_empty_code() -> None:
    with pytest.raises(CodeValidationError, match="No code was generated when generating tests"):
        validate_testgen_code("", python_version=(3, 11))


def test_whitespace_only() -> None:
    with pytest.raises(CodeValidationError, match="No code was generated when generating tests"):
        validate_testgen_code("   \n   \n", python_version=(3, 11))


def test_incomplete_function_definition() -> None:
    code: str = """def test_function():
    x = 1
    y = """

    # Individual validation removes the broken test, leaving no valid tests
    with pytest.raises(CodeValidationError, match="No test functions were found"):
        validate_testgen_code(code, python_version=(3, 11))


def test_single_broken_function() -> None:
    code: str = """def test_function():
        x =
        y = 2
        z = 3
        return x + y"""

    # A single broken test function is removed, leaving no valid tests
    with pytest.raises(CodeValidationError, match="No test functions were found"):
        validate_testgen_code(code, python_version=(3, 11))


def test_incomplete_statement_at_end() -> None:
    code: str = """def test_function():
    x = 1
    print("""

    # Individual validation removes the broken test, leaving no valid tests
    with pytest.raises(CodeValidationError, match="No test functions were found"):
        validate_testgen_code(code, python_version=(3, 11))


def test_syntax_error_in_middle() -> None:
    code: str = """def test_function():
    x = %
    y = 2
    return x + y"""

    # Individual validation removes the broken test, leaving no valid tests
    with pytest.raises(CodeValidationError, match="No test functions were found"):
        validate_testgen_code(code, python_version=(3, 11))


def test_multiple_functions_with_error() -> None:
    code: str = """def test_function1():
    return 1

def test_function2():
    x = 1
    return x

def test_function3():
    y = """

    expected: str = """def test_function1():
    return 1

def test_function2():
    x = 1
    return x"""

    # The broken test_function3 should be removed, valid tests kept
    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result == expected


def test_just_def() -> None:
    code: str = "def test_"
    with pytest.raises(CodeValidationError, match="Invalid test code generated"):
        validate_testgen_code(code, python_version=(3, 11))


def test_broken_def() -> None:
    code: str = """def test_function1():
    return 1

def test_function2():
    x = 1
    return x

def"""

    expected: str = """def test_function1():
    return 1

def test_function2():
    x = 1
    return x"""

    # The incomplete 'def' is not a test function, valid tests should be kept
    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result == expected


def test_pytest_style() -> None:
    code: str = """def test_something():
        pass"""
    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result.strip() == code.strip()


def test_unittest_style() -> None:
    code: str = """class TestCase(unittest.TestCase):
        def test_something(self):
            pass"""
    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result.strip() == code.strip()


def test_no_test_function() -> None:
    code: str = """def something():
        pass"""
    with pytest.raises(CodeValidationError, match="No test functions were found in the generated code"):
        validate_testgen_code(code, python_version=(3, 11))


def test_unittest_without_test_prefix() -> None:
    code: str = """class TestCase(unittest.TestCase):
        def something(self):
            pass"""
    with pytest.raises(CodeValidationError, match="No test functions were found in the generated code"):
        validate_testgen_code(code, python_version=(3, 11))


def test_remove_lines_with_test() -> None:
    code: str = """def test_one():
        pass

class MyTests(TestCase):
    def tes"""

    expected: str = """def test_one():
        pass"""

    # The incomplete class definition is removed, valid test kept
    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result == expected


def test_test_class_with_syntax_error() -> None:
    code: str = """import pytest

class TestValid:
    def test_one(self):
        pass

class TestBroken:
    def test_two(self):
        x = @invalid

class TestAlsoValid:
    def test_three(self):
        return 1"""

    expected: str = """import pytest


class TestValid:
    def test_one(self):
        pass

class TestAlsoValid:
    def test_three(self):
        return 1"""

    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result == expected


def test_existing_future_annotations() -> None:
    code: str = """from __future__ import annotations


def test_function():
    x: list[int] = [1, 2, 3]
    return x[0]"""

    expected: str = """from __future__ import annotations


def test_function():
    x: list[int] = [1, 2, 3]
    return x[0]"""

    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result == expected


def test_comments_before_future_import() -> None:
    code: str = """# This is a test file
import pytest as test

from __future__ import annotations

# With multiple comments
def test_function():
    x: list[int] = [1, 2, 3]
    return x[0]"""

    expected: str = """from __future__ import annotations

# This is a test file
import pytest as test


def test_function():
    x: list[int] = [1, 2, 3]
    return x[0]"""

    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result == expected


def test_multiple_future_imports() -> None:
    code: str = """from __future__ import division
from pytest import test
from __future__ import print_function

def test_function():
    x: list[int] = [1, 2, 3]
    return x[0]"""

    expected: str = """from __future__ import division, print_function

from pytest import test


def test_function():
    x: list[int] = [1, 2, 3]
    return x[0]"""

    result: str = validate_testgen_code(code, python_version=(3, 11))
    assert result == expected


def test_long_one_string() -> None:
    code = """
import pytest  # used for our unit tests
from dataclasses import dataclass
from typing import Any, TypeVar, Generic

# --- Dummy/Minimal implementations for required classes/types ---

class UserPromptPart:
    def __init__(self, content):
        self.content = content

class BinaryContent:
    def __init__(self, data: bytes, media_type: str):
        self.data = data
        self.media_type = media_type

class DocumentUrl:
    def __init__(self, url: str):
        self.url = url
        self.media_type = "application/pdf"

class AudioUrl:
    def __init__(self, url: str):
        self.url = url
        self.media_type = "audio/mpeg"

class ImageUrl:
    def __init__(self, url: str):
        self.url = url
        self.media_type = "image/png"

class VideoUrl:
    def __init__(self, url: str):
        self.url = url
        self.media_type = "video/mp4"

def assert_never(value):
    raise AssertionError(f"Unhandled value: {value}")

# --- The function to test, as given ---

WorkerOutputT = TypeVar('WorkerOutputT')
AgentDepsT = TypeVar('AgentDepsT')

@dataclass
class AgentWorker(Generic[WorkerOutputT, AgentDepsT]):
    agent: Any = None  # Not used in these tests

    def _request_parts_from_a2a(self, parts: list[dict]) -> list:
        model_parts: list = []
        for part in parts:
            if part['kind'] == 'text':
                model_parts.append(UserPromptPart(content=part['text']))
            elif part['kind'] == 'file':
                file_content = part['file']
                if 'bytes' in file_content:
                    data = file_content['bytes'].encode('utf-8')
                    mime_type = file_content.get('mime_type', 'application/octet-stream')
                    content = BinaryContent(data=data, media_type=mime_type)
                    model_parts.append(UserPromptPart(content=[content]))
                else:
                    url = file_content['uri']
                    for url_cls in (DocumentUrl, AudioUrl, ImageUrl, VideoUrl):
                        content = url_cls(url=url)
                        try:
                            content.media_type
                        except ValueError:
                            continue
                        else:
                            break
                    else:
                        raise ValueError(f'Unsupported file type: {url}')
                    model_parts.append(UserPromptPart(content=[content]))
            elif part['kind'] == 'data':
                raise NotImplementedError('Data parts are not supported yet.')
            else:
                assert_never(part)
        return model_parts

# --- Unit tests ---

@pytest.fixture
def worker():
    return AgentWorker()

# 1. Basic Test Cases

def test_single_text_part(worker):
    # Test conversion of a single text part
    parts = [{'kind': 'text', 'text': 'Hello world!'}]
    result = worker._request_parts_from_a2a(parts)
    assert len(result) == 1
    assert isinstance(result[0], UserPromptPart)
    assert result[0].content == 'Hello world!'

def test_single_file_part_bytes(worker):
    # Test conversion of a single file part with bytes content
    parts = [{'kind': 'file', 'file': {'bytes': 'abc123', 'mime_type': 'text/plain'}}]
    result = worker._request_parts_from_a2a(parts)
    assert len(result) == 1
    assert isinstance(result[0], UserPromptPart)
    content = result[0].content
    assert isinstance(content, list)
    assert isinstance(content[0], BinaryContent)
    assert content[0].data == b'abc123'
    assert content[0].media_type == 'text/plain'

def test_single_file_part_bytes_default_mime(worker):
    # Test file part with bytes but no mime_type (should default)
    parts = [{'kind': 'file', 'file': {'bytes': 'xyz'}}]
    result = worker._request_parts_from_a2a(parts)
    content = result[0].content[0]
    assert isinstance(content, BinaryContent)
    assert content.data == b'xyz'
    assert content.media_type == 'application/octet-stream'

def test_single_file_part_uri_document(worker):
    # Test file part with uri (should use DocumentUrl)
    parts = [{'kind': 'file', 'file': {'uri': 'http://example.com/doc.pdf'}}]
    result = worker._request_parts_from_a2a(parts)
    content = result[0].content[0]
    assert isinstance(content, DocumentUrl)
    assert content.url == 'http://example.com/doc.pdf'
    assert content.media_type == 'application/pdf'

def test_single_file_part_uri_audio(worker):
    # Test file part with uri (should use AudioUrl)
    # We'll simulate by patching the order so AudioUrl is picked
    # But in our implementation, DocumentUrl is always first, so we can't force AudioUrl without changing the code.
    # Instead, check that DocumentUrl is used, as per the code order.
    parts = [{'kind': 'file', 'file': {'uri': 'http://example.com/audio.mp3'}}]
    result = worker._request_parts_from_a2a(parts)
    content = result[0].content[0]
    assert isinstance(content, DocumentUrl)
    assert content.url == 'http://example.com/audio.mp3'
    assert content.media_type == 'application/pdf'

def test_multiple_parts_mixed(worker):
    # Test multiple parts: text, file (bytes), file (uri)
    parts = [
        {'kind': 'text', 'text': 'A'},
        {'kind': 'file', 'file': {'bytes': 'bin', 'mime_type': 'application/octet-stream'}},
        {'kind': 'file', 'file': {'uri': 'http://example.com/img.png'}}
    ]
    result = worker._request_parts_from_a2a(parts)
    assert len(result) == 3
    assert isinstance(result[0], UserPromptPart)
    assert result[0].content == 'A'
    assert isinstance(result[1].content[0], BinaryContent)
    assert result[1].content[0].data == b'bin'
    assert result[1].content[0].media_type == 'application/octet-stream'
    assert isinstance(result[2].content[0], DocumentUrl)
    assert result[2].content[0].url == 'http://example.com/img.png'

# 2. Edge Test Cases

def test_empty_parts_list(worker):
    # Test empty input list
    parts = []
    result = worker._request_parts_from_a2a(parts)
    assert result == []

def test_unsupported_kind(worker):
    # Test with an unsupported kind
    parts = [{'kind': 'foo', 'foo': 'bar'}]
    with pytest.raises(AssertionError):
        worker._request_parts_from_a2a(parts)

def test_file_part_missing_bytes_and_uri(worker):
    # Test file part with neither 'bytes' nor 'uri'
    parts = [{'kind': 'file', 'file': {}}]
    with pytest.raises(KeyError):
        worker._request_parts_from_a2a(parts)

def test_file_part_with_unknown_url(worker):
    # Test file part with a uri that doesn't match any url class (simulate by making all url classes fail)
    # In our code, DocumentUrl always works, so this is not possible unless we patch DocumentUrl.
    # We'll skip this as our implementation always picks DocumentUrl.

def test_data_kind_not_implemented(worker):
    # Test with 'data' kind, which should raise NotImplementedError
    parts = [{'kind': 'data', 'data': 'foo'}]
    with pytest.raises(NotImplementedError):
        worker._request_parts_from_a2a(parts)

def test_text_part_missing_text_key(worker):
    # Test text part missing the 'text' key
    parts = [{'kind': 'text'}]
    with pytest.raises(KeyError):
        worker._request_parts_from_a2a(parts)

def test_file_part_bytes_non_utf8(worker):
    # Test file part with bytes that aren't valid utf-8 (simulate with a string that would fail if decoded, but .encode never fails)
    # Since we use .encode, not decode, any string is accepted. If bytes were passed instead, it would fail.
    # So, let's pass a bytes object instead of string to simulate a user error:
    parts = [{'kind': 'file', 'file': {'bytes': b'\xff\xfe'}}]
    # This will fail because b'\xff\xfe'.encode('utf-8') is not valid (bytes has no encode)
    with pytest.raises(AttributeError):
        worker._request_parts_from_a2a(parts)

def test_file_part_bytes_empty_string(worker):
    # Test file part with empty bytes string
    parts = [{'kind': 'file', 'file': {'bytes': '', 'mime_type': 'text/plain'}}]
    result = worker._request_parts_from_a2a(parts)
    content = result[0].content[0]
    assert content.data == b''
    assert content.media_...
"""

    assert validate_testgen_code(code, python_version=(3, 11))


class TestDecoratorHandling:
    """Tests for decorator preservation during code splitting."""

    def test_split_code_with_ast_preserves_decorators(self) -> None:
        """AST-based splitting should include decorators with test functions."""
        code = """import pytest

@pytest.mark.asyncio
async def test_foo():
    pass

@pytest.fixture
def my_fixture():
    return 1
"""
        result = split_code_with_ast(code, (3, 11))
        assert result is not None
        assert len(result.test_functions) == 1

        expected_test = """@pytest.mark.asyncio
async def test_foo():
    pass
"""
        assert result.test_functions[0] == expected_test

        expected_preamble = """import pytest

@pytest.fixture
def my_fixture():
    return 1
"""
        assert result.preamble == expected_preamble

    def test_split_code_with_ast_preserves_multiple_decorators(self) -> None:
        """AST-based splitting should preserve multiple stacked decorators."""
        code = """import pytest

@pytest.mark.asyncio
@pytest.mark.slow
@pytest.mark.integration
async def test_complex():
    pass
"""
        result = split_code_with_ast(code, (3, 11))
        assert result is not None
        assert len(result.test_functions) == 1

        expected_test = """@pytest.mark.asyncio
@pytest.mark.slow
@pytest.mark.integration
async def test_complex():
    pass
"""
        assert result.test_functions[0] == expected_test

    def test_split_code_with_regex_preserves_decorators(self) -> None:
        """Regex-based splitting should include decorators with test functions."""
        code = """import pytest

@pytest.mark.asyncio
async def test_foo():
    pass

def test_bar():
    pass
"""
        result = split_code_with_regex(code)
        assert len(result.test_functions) == 2

        expected_first = """@pytest.mark.asyncio
async def test_foo():
    pass"""
        assert result.test_functions[0] == expected_first

        expected_second = """def test_bar():
    pass"""
        assert result.test_functions[1] == expected_second

    def test_split_code_with_regex_preserves_multiple_decorators(self) -> None:
        """Regex-based splitting should preserve multiple stacked decorators."""
        code = """import pytest

@pytest.mark.asyncio
@pytest.mark.slow
async def test_foo():
    pass
"""
        result = split_code_with_regex(code)
        assert len(result.test_functions) == 1

        expected_test = """@pytest.mark.asyncio
@pytest.mark.slow
async def test_foo():
    pass"""
        assert result.test_functions[0] == expected_test

    def test_validate_testgen_code_preserves_asyncio_decorator(self) -> None:
        """End-to-end test: asyncio decorators should be preserved in final output."""
        code = """import pytest

@pytest.mark.asyncio
async def test_async_function():
    await something()
    pass
"""
        expected = """import pytest


@pytest.mark.asyncio
async def test_async_function():
    await something()
    pass"""

        result = validate_testgen_code(code, python_version=(3, 11))
        assert result == expected

    def test_validate_testgen_code_preserves_parametrize_decorator(self) -> None:
        """Parametrize decorators should be preserved in final output."""
        code = """import pytest

@pytest.mark.parametrize("x,y", [(1, 2), (3, 4)])
def test_parametrized(x, y):
    assert x < y
"""
        expected = """import pytest


@pytest.mark.parametrize("x,y", [(1, 2), (3, 4)])
def test_parametrized(x, y):
    assert x < y"""

        result = validate_testgen_code(code, python_version=(3, 11))
        assert result == expected

    def test_multiple_decorated_tests_preserved(self) -> None:
        """Multiple tests each with decorators should all be preserved."""
        code = """import pytest

@pytest.mark.asyncio
async def test_first():
    pass

@pytest.mark.slow
def test_second():
    pass

@pytest.mark.asyncio
@pytest.mark.integration
async def test_third():
    pass
"""
        expected = """import pytest


@pytest.mark.asyncio
async def test_first():
    pass


@pytest.mark.slow
def test_second():
    pass


@pytest.mark.asyncio
@pytest.mark.integration
async def test_third():
    pass"""

        result = validate_testgen_code(code, python_version=(3, 11))
        assert result == expected


class TestRepairPreamble:
    """Tests for preamble repair when LLM generates broken import/helper code."""

    def test_repair_removes_single_broken_line(self) -> None:
        preamble = """import pytest
x = @invalid
import os"""
        result = repair_preamble(preamble, (3, 11))
        assert "import pytest" in result
        assert "import os" in result
        assert "@invalid" not in result

    def test_repair_returns_empty_for_all_broken(self) -> None:
        preamble = "x = @\ny = %\nz = $"
        result = repair_preamble(preamble, (3, 11))
        assert result == ""

    def test_repair_noop_for_valid_preamble(self) -> None:
        preamble = "import pytest\nimport os"
        result = repair_preamble(preamble, (3, 11))
        assert "import pytest" in result
        assert "import os" in result

    def test_repair_handles_empty_string(self) -> None:
        assert repair_preamble("", (3, 11)) == ""

    def test_repair_removes_truncated_string(self) -> None:
        """LLMs sometimes produce unclosed strings in helper code."""
        preamble = """import pytest
EXPECTED = "hello world
import os"""
        result = repair_preamble(preamble, (3, 11))
        # The broken string line should be removed
        assert "import pytest" in result
        assert "hello world" not in result

    def test_repair_removes_incomplete_function(self) -> None:
        preamble = """import pytest

def helper():
    x = 1

def broken_helper(
    # missing closing paren"""
        result = repair_preamble(preamble, (3, 11))
        assert "import pytest" in result
        assert "def helper" in result


class TestBrokenPreambleValidation:
    """End-to-end tests: broken preamble should not discard valid tests."""

    def test_broken_preamble_keeps_valid_tests(self) -> None:
        """When preamble has a syntax error, valid tests should still be kept."""
        code = """import pytest
x = @invalid_syntax

def test_one():
    assert 1 == 1

def test_two():
    assert 2 == 2"""

        result = validate_testgen_code(code, python_version=(3, 11))
        assert "def test_one" in result
        assert "def test_two" in result

    def test_broken_preamble_with_truncated_string(self) -> None:
        """Truncated string in preamble should not kill all tests."""
        code = """import pytest
EXPECTED = "unterminated

def test_basic():
    assert True"""

        result = validate_testgen_code(code, python_version=(3, 11))
        assert "def test_basic" in result

    def test_broken_preamble_mixed_valid_invalid_tests(self) -> None:
        """Broken preamble + some broken tests should keep only the valid tests."""
        code = """import pytest
x = @bad

def test_valid():
    assert True

def test_broken():
    y = @also_bad"""

        result = validate_testgen_code(code, python_version=(3, 11))
        assert "def test_valid" in result
        assert "test_broken" not in result
