"""Tests for parse_and_validate_llm_output function.

These tests verify that function definitions are replaced with imports
before validation, fixing the case where LLMs include the entire function
definition instead of importing it.
"""

from unittest.mock import MagicMock

import pytest

from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize
from core.languages.python.testgen.testgen import parse_and_validate_llm_output


@pytest.fixture
def mock_context() -> MagicMock:
    """Create a mock context that doesn't trigger ellipsis errors."""
    ctx = MagicMock()
    ctx.did_generate_ellipsis.return_value = False
    return ctx


def test_replaces_function_definition_with_import_before_validation(mock_context: MagicMock) -> None:
    """Test that function definitions are replaced with imports before validation.

    This is the specific case that failed: the LLM generated test code that included
    the entire function definition (~400 lines) instead of importing it. Without
    early replacement, validation would fail on such large preambles.
    """
    # Simulate LLM response with function definition instead of import
    llm_response = '''Here are the tests:

```python
import pytest

# The LLM incorrectly included the function definition instead of importing it
def comparator(orig, new, superset_obj=False):
    """Compare two objects for equality recursively."""
    if type(orig) is not type(new):
        return False
    if isinstance(orig, (list, tuple)):
        if len(orig) != len(new):
            return False
        return all(comparator(e1, e2, superset_obj) for e1, e2 in zip(orig, new))
    if isinstance(orig, dict):
        if len(orig) != len(new):
            return False
        return all(k in new and comparator(v, new[k], superset_obj) for k, v in orig.items())
    return orig == new


def test_basic_equality():
    assert comparator(1, 1) is True
    assert comparator(1, 2) is False


def test_list_comparison():
    assert comparator([1, 2, 3], [1, 2, 3]) is True
    assert comparator([1, 2], [1, 2, 3]) is False


def test_dict_comparison():
    assert comparator({"a": 1}, {"a": 1}) is True
    assert comparator({"a": 1}, {"a": 2}) is False
```
'''

    function_to_optimize = FunctionToOptimize(
        function_name="comparator", file_path="/path/to/codeflash/verification/comparator.py", parents=[]
    )
    module_path = "codeflash.verification.comparator"

    result = parse_and_validate_llm_output(
        response_content=llm_response,
        ctx=mock_context,
        python_version=(3, 11, 0),
        error_context="",
        function_to_optimize=function_to_optimize,
        module_path=module_path,
    )

    expected = """import pytest
from codeflash.verification.comparator import comparator


def test_basic_equality():
    assert comparator(1, 1) is True
    assert comparator(1, 2) is False


def test_list_comparison():
    assert comparator([1, 2, 3], [1, 2, 3]) is True
    assert comparator([1, 2], [1, 2, 3]) is False


def test_dict_comparison():
    assert comparator({"a": 1}, {"a": 1}) is True
    assert comparator({"a": 1}, {"a": 2}) is False"""

    assert result.strip() == expected.strip()


def test_preserves_test_code_when_no_function_to_replace(mock_context: MagicMock) -> None:
    """Test that code is preserved when there's no function definition to replace."""
    llm_response = """```python
import pytest
from mymodule import my_function


def test_basic():
    assert my_function(1) == 2
```
"""

    function_to_optimize = FunctionToOptimize(function_name="my_function", file_path="/path/to/mymodule.py", parents=[])
    module_path = "mymodule"

    result = parse_and_validate_llm_output(
        response_content=llm_response,
        ctx=mock_context,
        python_version=(3, 11, 0),
        error_context="",
        function_to_optimize=function_to_optimize,
        module_path=module_path,
    )

    expected = """import pytest
from mymodule import my_function


def test_basic():
    assert my_function(1) == 2"""

    assert result.strip() == expected.strip()


def test_works_without_function_to_optimize(mock_context: MagicMock) -> None:
    """Test that validation still works when function_to_optimize is None."""
    llm_response = """```python
import pytest


def test_simple():
    assert 1 == 1
```
"""

    result = parse_and_validate_llm_output(
        response_content=llm_response,
        ctx=mock_context,
        python_version=(3, 11, 0),
        error_context="",
        function_to_optimize=None,
        module_path=None,
    )

    expected = """import pytest


def test_simple():
    assert 1 == 1"""

    assert result.strip() == expected.strip()


def test_handles_large_function_definition(mock_context: MagicMock) -> None:
    """Test that large function definitions (like the ~400 line comparator) are handled.

    This simulates the actual failure case where the LLM included a very large
    function definition that might cause validation issues.
    """
    # Create a function with multiple lines (simulating a larger function)
    llm_response = '''```python
import pytest

def large_function(arg):
    """A large function that should be replaced."""
    x0 = 0
    x1 = 1
    x2 = 2
    x3 = 3
    x4 = 4
    return arg


def test_large_function():
    assert large_function(1) == 1


def test_large_function_edge_case():
    assert large_function(0) == 0
```
'''

    function_to_optimize = FunctionToOptimize(
        function_name="large_function", file_path="/path/to/module.py", parents=[]
    )
    module_path = "myapp.utils"

    result = parse_and_validate_llm_output(
        response_content=llm_response,
        ctx=mock_context,
        python_version=(3, 11, 0),
        error_context="",
        function_to_optimize=function_to_optimize,
        module_path=module_path,
    )

    expected = """import pytest
from myapp.utils import large_function


def test_large_function():
    assert large_function(1) == 1


def test_large_function_edge_case():
    assert large_function(0) == 0"""

    assert result.strip() == expected.strip()


def test_replaces_class_definition_with_import(mock_context: MagicMock) -> None:
    """Test that class definitions can also be replaced with imports."""
    llm_response = """```python
import pytest

class MyClass:
    def __init__(self, value):
        self.value = value

    def get_value(self):
        return self.value


def test_my_class():
    obj = MyClass(42)
    assert obj.get_value() == 42
```
"""

    function_to_optimize = FunctionToOptimize(
        function_name="MyClass",
        file_path="/path/to/module.py",
        parents=[FunctionParent(name="MyClass", type="ClassDef")],
    )
    module_path = "myapp.models"

    result = parse_and_validate_llm_output(
        response_content=llm_response,
        ctx=mock_context,
        python_version=(3, 11, 0),
        error_context="",
        function_to_optimize=function_to_optimize,
        module_path=module_path,
    )

    expected = """import pytest
from myapp.models import MyClass


def test_my_class():
    obj = MyClass(42)
    assert obj.get_value() == 42"""

    assert result.strip() == expected.strip()
