"""Tests for unused helper function revert functionality."""

from __future__ import annotations

from pathlib import Path

import pytest

from codeflash_python._model import FunctionParent, FunctionToOptimize
from codeflash_python.context.models import CodeStringsMarkdown
from codeflash_python.context.pipeline import get_code_optimization_context
from codeflash_python.pipeline._function_optimizer import apply_optimized_code
from codeflash_python.verification._unused_helpers import (
    detect_unused_helper_functions,
    revert_unused_helper_functions,
)


@pytest.fixture
def temp_project(tmp_path: Path) -> tuple[Path, Path]:
    """Create a temporary project with test files."""
    temp_dir = tmp_path

    # Main file with function that calls helpers
    main_file = temp_dir / "main.py"
    main_file.write_text("""
def entrypoint_function(n):
    \"\"\"Function that calls two helper functions.\"\"\"
    result1 = helper_function_1(n)
    result2 = helper_function_2(n)
    return result1 + result2

def helper_function_1(x):
    \"\"\"First helper function.\"\"\"
    return x * 2

def helper_function_2(x):
    \"\"\"Second helper function.\"\"\"
    return x * 3
""")

    return temp_dir, main_file


def test_detect_unused_helper_functions(
    temp_project: tuple[Path, Path],
) -> None:
    """Test that unused helper functions are correctly detected."""
    temp_dir, main_file = temp_project

    # Optimized version that only calls one helper
    optimized_code = """
```python:main.py
def entrypoint_function(n):
    \"\"\"Optimized function that only calls one helper.\"\"\"
    result1 = helper_function_1(n)
    return result1 + n * 3  # Inlined helper_function_2

def helper_function_1(x):
    \"\"\"First helper function.\"\"\"
    return x * 2

def helper_function_2(x):
    \"\"\"Second helper function - MODIFIED VERSION should be reverted.\"\"\"
    return x * 4  # This change should be reverted to original x * 3
```
"""

    # Create FunctionToOptimize instance
    function_to_optimize = FunctionToOptimize(
        file_path=main_file, function_name="entrypoint_function", parents=()
    )

    # Get original code context to find helper functions
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect helper_function_2 as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    expected_unused = {"helper_function_2"}

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )

    # Also test the complete replace workflow
    # First modify the optimized code to include a MODIFIED unused helper
    optimized_code_with_modified_helper = """
```python:main.py
def entrypoint_function(n):
    \"\"\"Optimized function that only calls one helper.\"\"\"
    result1 = helper_function_1(n)
    return result1 + n * 3  # Inlined helper_function_2

def helper_function_1(x):
    \"\"\"First helper function.\"\"\"
    return x * 2

def helper_function_2(x):
    \"\"\"Second helper function - MODIFIED VERSION should be reverted.\"\"\"
    return x * 7  # This should be reverted to x * 3
```
"""

    original_helper_code = {main_file: main_file.read_text()}

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code_with_modified_helper,
        original_helper_code,
        temp_dir,
    )
    # Check final file content
    final_content = main_file.read_text()

    # The entrypoint should be optimized
    assert "result1 + n * 3" in final_content, (
        "Entrypoint function should be optimized"
    )

    # helper_function_2 should be reverted to original (x * 3, NOT the modified x * 7)
    assert "return x * 3" in final_content, (
        "helper_function_2 should be reverted to original"
    )
    assert "return x * 7" not in final_content, (
        "helper_function_2 should NOT contain the modified version"
    )

    # helper_function_1 should remain (it's still called)
    assert "def helper_function_1(x):" in final_content, (
        "helper_function_1 should still exist"
    )

    # Also test the complete replace workflow
    original_helper_code = {main_file: main_file.read_text()}

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )

    # Check final file content
    final_content = main_file.read_text()

    # The entrypoint should be optimized
    assert "result1 + n * 3" in final_content, (
        "Entrypoint function should be optimized"
    )

    # helper_function_2 should be reverted to original (return x * 3, NOT the modified x * 4)
    assert "return x * 3" in final_content, (
        "helper_function_2 should be reverted to original"
    )
    assert "return x * 4" not in final_content, (
        "helper_function_2 should NOT contain the modified version"
    )

    # helper_function_1 should remain as optimized (it's still called)
    assert "def helper_function_1(x):" in final_content, (
        "helper_function_1 should still exist"
    )


def test_revert_unused_helper_functions(
    temp_project: tuple[Path, Path],
) -> None:
    """Test that unused helper functions are correctly reverted to original definitions."""
    temp_dir, main_file = temp_project

    # Optimized version that only calls one helper and modifies the unused one
    optimized_code = """
```python:main.py
def entrypoint_function(n):
    \"\"\"Optimized function that only calls one helper.\"\"\"
    result1 = helper_function_1(n)
    return result1 + n * 3  # Inlined helper_function_2

def helper_function_1(x):
    \"\"\"First helper function.\"\"\"
    return x * 2

def helper_function_2(x):
    \"\"\"Modified helper function - should be reverted.\"\"\"
    return x * 4  # This change should be reverted
```
"""

    # Create FunctionToOptimize instance
    function_to_optimize = FunctionToOptimize(
        file_path=main_file, function_name="entrypoint_function", parents=()
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Store original helper code
    original_helper_code = {main_file: main_file.read_text()}
    original_content = main_file.read_text()

    # Test the new functionality - this should:
    # 1. Apply the optimization
    # 2. Detect unused helpers
    # 3. Revert unused helpers to original definitions
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )

    # Check final file content
    final_content = main_file.read_text()

    # The entrypoint should be optimized (inline the helper_function_2 call)
    assert "result1 + n * 3" in final_content, (
        "Entrypoint function should be optimized"
    )

    # helper_function_2 should be reverted to original (return x * 3, not x * 4)
    assert "return x * 3" in final_content, (
        "helper_function_2 should be reverted to original"
    )
    assert "return x * 4" not in final_content, (
        "helper_function_2 should not contain the optimized version"
    )

    # helper_function_1 should remain as optimized (it's still called)
    assert "def helper_function_1(x):" in final_content, (
        "helper_function_1 should still exist"
    )


def test_no_unused_helpers_no_revert(temp_project: tuple[Path, Path]) -> None:
    """Test that when all helpers are still used, nothing is reverted."""
    temp_dir, main_file = temp_project

    # Store original content to verify nothing changes
    original_content = main_file.read_text()

    revert_unused_helper_functions(temp_dir, [], {})

    # Verify the file content remains unchanged
    assert main_file.read_text() == original_content, (
        "File should remain unchanged when no helpers to revert"
    )

    # Optimized version that still calls both helpers
    optimized_code = """
```python:main.py
def entrypoint_function(n):
    \"\"\"Optimized function that still calls both helpers.\"\"\"
    result1 = helper_function_1(n)
    result2 = helper_function_2(n)
    return result1 + result2  # Still using both

def helper_function_1(x):
    \"\"\"First helper function - optimized.\"\"\"
    return x << 1  # Optimized to use bit shift

def helper_function_2(x):
    \"\"\"Second helper function - optimized.\"\"\"
    return x * 3
```
"""

    # Create FunctionToOptimize instance
    function_to_optimize = FunctionToOptimize(
        file_path=main_file, function_name="entrypoint_function", parents=()
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Store original helper code
    original_helper_code = {main_file: main_file.read_text()}

    # Test detection - should find no unused helpers
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )
    assert len(unused_helpers) == 0, "No helpers should be detected as unused"

    # Apply optimization
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )

    # Check final file content - should contain the optimized versions
    final_content = main_file.read_text()

    # Both helpers should be optimized
    assert "x << 1" in final_content, (
        "helper_function_1 should be optimized to use bit shift"
    )
    assert "result1 + result2" in final_content, (
        "Entrypoint should still call both helpers"
    )


def test_detect_unused_in_multi_file_project(tmp_path: Path) -> None:
    """Test detection of unused helpers across multiple files."""
    temp_dir = tmp_path

    # Main file
    main_file = temp_dir / "main.py"
    main_file.write_text("""
from helpers import helper_function_1, helper_function_2

def entrypoint_function(n):
    \"\"\"Function that calls helpers from another file.\"\"\"
    result1 = helper_function_1(n)
    result2 = helper_function_2(n)
    return result1 + result2
""")

    # Helper file
    helper_file = temp_dir / "helpers.py"
    helper_file.write_text("""
def helper_function_1(x):
    \"\"\"First helper function.\"\"\"
    return x * 2

def helper_function_2(x):
    \"\"\"Second helper function.\"\"\"
    return x * 3

def helper_function_1(y):  # Duplicate name to test line 575
    \"\"\"Overloaded helper function.\"\"\"
    return y + 10
""")

    # Optimized version that only calls one helper with aliased import
    optimized_code = """
```python:main.py
from helpers import helper_function_1 as h1
import helpers as h_module

def entrypoint_function(n):
    \"\"\"Optimized function that only calls one helper with aliasing.\"\"\"
    result1 = h1(n)  # Using aliased import
    # Inlined helper_function_2 functionality: n * 3
    return result1 + n * 3  # Fully inlined helper_function_2
```
"""

    # Create FunctionToOptimize instance
    function_to_optimize = FunctionToOptimize(
        file_path=main_file, function_name="entrypoint_function", parents=()
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect helper_function_2 as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    expected_unused = {"helper_function_2"}

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )

    # Also test the complete replace workflow
    # First, simulate modified helper in the helper file
    helper_file.write_text("""
def helper_function_1(x):
    \"\"\"First helper function.\"\"\"
    return x * 2

def helper_function_2(x):
    \"\"\"Second helper function - MODIFIED VERSION.\"\"\"
    return x * 9  # This should be reverted to x * 3
""")

    # Store original helper code (before modification)
    original_helper_code = {
        main_file: """
from helpers import helper_function_1, helper_function_2

def entrypoint_function(n):
    \"\"\"Function that calls helpers from another file.\"\"\"
    result1 = helper_function_1(n)
    result2 = helper_function_2(n)
    return result1 + result2
""",
        helper_file: """
def helper_function_1(x):
    \"\"\"First helper function.\"\"\"
    return x * 2

def helper_function_2(x):
    \"\"\"Second helper function.\"\"\"
    return x * 3
""",
    }

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )
    # Check main file content
    main_content = main_file.read_text()
    assert "result1 + n * 3" in main_content, (
        "Entrypoint function should be optimized"
    )
    assert "from helpers import helper_function_1" in main_content, (
        "Import should be updated"
    )

    # Check helper file content - helper_function_2 should be reverted to original
    helper_content = helper_file.read_text()
    assert "def helper_function_1(x):" in helper_content, (
        "helper_function_1 should still exist"
    )
    assert "def helper_function_2(x):" in helper_content, (
        "helper_function_2 should exist"
    )
    assert "return x * 3" in helper_content, (
        "helper_function_2 should be reverted to original"
    )
    assert "return x * 9" not in helper_content, (
        "helper_function_2 should NOT contain the modified version"
    )

    # Also test the complete replace workflow
    # First, simulate modified helper in the helper file
    helper_file.write_text("""
def helper_function_1(x):
    \"\"\"First helper function.\"\"\"
    return x * 2

def helper_function_2(x):
    \"\"\"Second helper function - MODIFIED VERSION.\"\"\"
    return x * 5  # This should be reverted to x * 3
""")

    # Store original helper code (before modification)
    original_helper_code = {
        main_file: """
from helpers import helper_function_1, helper_function_2

def entrypoint_function(n):
    \"\"\"Function that calls helpers from another file.\"\"\"
    result1 = helper_function_1(n)
    result2 = helper_function_2(n)
    return result1 + result2
""",
        helper_file: """
def helper_function_1(x):
    \"\"\"First helper function.\"\"\"
    return x * 2

def helper_function_2(x):
    \"\"\"Second helper function.\"\"\"
    return x * 3
""",
    }

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )

    # Check main file content
    main_content = main_file.read_text()
    assert "result1 + n * 3" in main_content, (
        "Entrypoint function should be optimized"
    )
    assert "from helpers import helper_function_1" in main_content, (
        "Import should be updated"
    )

    # Check helper file content - helper_function_2 should be reverted to original
    helper_content = helper_file.read_text()
    assert "def helper_function_1(x):" in helper_content, (
        "helper_function_1 should still exist"
    )
    assert "def helper_function_2(x):" in helper_content, (
        "helper_function_2 should exist"
    )
    assert "return x * 3" in helper_content, (
        "helper_function_2 should be reverted to original"
    )
    assert "return x * 5" not in helper_content, (
        "helper_function_2 should NOT contain the modified version"
    )


def test_class_method_entrypoint_with_helper_methods(tmp_path: Path) -> None:
    """Test unused helper detection when entrypoint is a class method that calls other methods."""
    temp_dir = tmp_path

    # Main file with class containing methods
    main_file = temp_dir / "main.py"
    main_file.write_text("""
class Calculator:
    def entrypoint_method(self, n):
        \"\"\"Main method that calls helper methods.\"\"\"
        result1 = self.helper_method_1(n)
        result2 = self.helper_method_2(n)
        return result1 + result2

    def helper_method_1(self, x):
        \"\"\"First helper method.\"\"\"
        return x * 2

    def helper_method_2(self, x):
        \"\"\"Second helper method.\"\"\"
        return x * 3
""")

    # Optimized version that only calls one helper method
    optimized_code = """
```python:main.py
class Calculator:
    def entrypoint_method(self, n):
        \"\"\"Optimized method that only calls one helper.\"\"\"
        result1 = self.helper_method_1(n)
        return result1 + n * 3  # Inlined helper_method_2

    def helper_method_1(self, x):
        \"\"\"First helper method.\"\"\"
        return x * 2

    def helper_method_2(self, x):
        \"\"\"Second helper method - should be reverted.\"\"\"
        return x * 4
```
"""

    # Create FunctionToOptimize instance for class method
    function_to_optimize = FunctionToOptimize(
        file_path=main_file,
        function_name="entrypoint_method",
        parents=(FunctionParent(name="Calculator", type="ClassDef"),),
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect Calculator.helper_method_2 as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    # The new context pipeline includes the entrypoint itself in helper_functions,
    # so it also appears as "unused" (it doesn't call itself).
    expected_unused = {
        "Calculator.helper_method_2",
        "Calculator.entrypoint_method",
    }

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )

    # Also test the complete replace workflow
    # Update optimized code to include a MODIFIED unused helper
    optimized_code_with_modified_helper = """
```python:main.py
class Calculator:
    def entrypoint_method(self, n):
        \"\"\"Optimized method that only calls one helper.\"\"\"
        result1 = self.helper_method_1(n)
        return result1 + n * 3  # Inlined helper_method_2

    def helper_method_1(self, x):
        \"\"\"First helper method.\"\"\"
        return x * 2

    def helper_method_2(self, x):
        \"\"\"Second helper method - MODIFIED VERSION should be reverted.\"\"\"
        return x * 8  # This should be reverted to x * 3
```
"""

    original_helper_code = {main_file: main_file.read_text()}

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code_with_modified_helper,
        original_helper_code,
        temp_dir,
    )

    # Check final file content
    final_content = main_file.read_text()

    # The entrypoint method should be optimized
    assert "result1 + n * 3" in final_content, (
        "Entrypoint method should be optimized"
    )

    # helper_method_2 should be reverted to original (x * 3, NOT the modified x * 8)
    assert "return x * 3" in final_content, (
        "helper_method_2 should be reverted to original"
    )
    assert "return x * 8" not in final_content, (
        "helper_method_2 should NOT contain the modified version"
    )

    # helper_method_1 should remain (it's still called)
    assert "def helper_method_1(self, x):" in final_content, (
        "helper_method_1 should still exist"
    )

    # Test reversion
    original_helper_code = {main_file: main_file.read_text()}

    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )

    # Check final file content
    final_content = main_file.read_text()

    # The entrypoint method should be optimized
    assert "result1 + n * 3" in final_content, (
        "Entrypoint method should be optimized"
    )

    # helper_method_2 should be reverted to original
    assert "x * 3" in final_content, "helper_method_2 should still exist"


def test_class_method_calls_external_helper_functions(tmp_path: Path) -> None:
    """Test when class method calls external helper functions."""
    temp_dir = tmp_path

    # Main file with class method that calls external helpers
    main_file = temp_dir / "main.py"
    main_file.write_text("""
def external_helper_1(x):
    \"\"\"External helper function.\"\"\"
    return x * 2

def external_helper_2(x):
    \"\"\"External helper function.\"\"\"
    return x * 3

class Processor:
    def process_data(self, n):
        \"\"\"Method that calls external helper functions.\"\"\"
        result1 = external_helper_1(n)
        result2 = external_helper_2(n)
        return result1 + result2
""")

    # Optimized version that only calls one external helper
    optimized_code = """
```python:main.py
def external_helper_1(x):
    \"\"\"External helper function.\"\"\"
    return x * 2

def external_helper_2(x):
    \"\"\"External helper function - should be reverted.\"\"\"
    return x * 3

class Processor:
    def process_data(self, n):
        \"\"\"Optimized method that only calls one helper.\"\"\"
        result1 = external_helper_1(n)
        return result1 + n * 3  # Inlined external_helper_2
```
"""

    # Create FunctionToOptimize instance for class method
    function_to_optimize = FunctionToOptimize(
        file_path=main_file,
        function_name="process_data",
        parents=(FunctionParent(name="Processor", type="ClassDef"),),
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect external_helper_2 as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    # The new context pipeline includes the entrypoint itself in helper_functions,
    # so it also appears as "unused" (it doesn't call itself).
    expected_unused = {"external_helper_2", "Processor.process_data"}

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )

    # Also test the complete replace workflow
    # Update optimized code to include a MODIFIED unused helper
    optimized_code_with_modified_helper = """
```python:main.py
def external_helper_1(x):
    \"\"\"External helper function.\"\"\"
    return x * 2

def external_helper_2(x):
    \"\"\"External helper function - MODIFIED VERSION should be reverted.\"\"\"
    return x * 11  # This should be reverted to x * 3

class Processor:
    def process_data(self, n):
        \"\"\"Optimized method that only calls one helper.\"\"\"
        result1 = external_helper_1(n)
        return result1 + n * 3  # Inlined external_helper_2
```
"""

    original_helper_code = {main_file: main_file.read_text()}

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code_with_modified_helper,
        original_helper_code,
        temp_dir,
    )

    # Check final file content
    final_content = main_file.read_text()

    # The class method should be optimized
    assert "result1 + n * 3" in final_content, (
        "Process method should be optimized"
    )

    # external_helper_2 should be reverted to original (x * 3, NOT the modified x * 11)
    assert "return x * 3" in final_content, (
        "external_helper_2 should be reverted to original"
    )
    assert "return x * 11" not in final_content, (
        "external_helper_2 should NOT contain the modified version"
    )

    # external_helper_1 should remain (it's still called)
    assert "def external_helper_1(x):" in final_content, (
        "external_helper_1 should still exist"
    )

    # Also test the complete replace workflow
    # Update optimized code to include a MODIFIED unused helper
    optimized_code_with_modified_helper = """
```python:main.py
def external_helper_1(x):
    \"\"\"External helper function.\"\"\"
    return x * 2

def external_helper_2(x):
    \"\"\"External helper function - MODIFIED VERSION should be reverted.\"\"\"
    return x * 7  # This should be reverted to x * 3

class Processor:
    def process_data(self, n):
        \"\"\"Optimized method that only calls one helper.\"\"\"
        result1 = external_helper_1(n)
        return result1 + n * 3  # Inlined external_helper_2
```
"""

    original_helper_code = {main_file: main_file.read_text()}

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code_with_modified_helper,
        original_helper_code,
        temp_dir,
    )

    # Check final file content
    final_content = main_file.read_text()

    # The class method should be optimized
    assert "result1 + n * 3" in final_content, (
        "Process method should be optimized"
    )

    # external_helper_2 should be reverted to original (x * 3, NOT the modified x * 7)
    assert "return x * 3" in final_content, (
        "external_helper_2 should be reverted to original"
    )
    assert "return x * 7" not in final_content, (
        "external_helper_2 should NOT contain the modified version"
    )

    # external_helper_1 should remain (it's still called)
    assert "def external_helper_1(x):" in final_content, (
        "external_helper_1 should still exist"
    )


def test_nested_class_method_optimization(tmp_path: Path) -> None:
    """Test optimization of methods in nested classes."""
    temp_dir = tmp_path

    # Main file with nested class
    main_file = temp_dir / "main.py"
    main_file.write_text("""
def global_helper_1(x):
    return x * 2

def global_helper_2(x):
    return x * 3

class OuterClass:
    class InnerProcessor:
        def compute(self, n):
            \"\"\"Method that calls global helper functions.\"\"\"
            result1 = global_helper_1(n)
            result2 = global_helper_2(n)
            return result1 + result2

        def local_helper(self, x):
            return x + 1
""")

    # Optimized version that inlines one helper
    optimized_code = """
```python:main.py
def global_helper_1(x):
    return x * 2

def global_helper_2(x):
    return x * 3

class OuterClass:
    class InnerProcessor:
        def compute(self, n):
            \"\"\"Optimized method.\"\"\"
            result1 = global_helper_1(n)
            return result1 + n * 3  # Inlined global_helper_2

        def local_helper(self, x):
            return x + 1
```
"""

    # Note: In practice, codeflash might not handle deeply nested classes,
    # but we test the detection logic anyway
    function_to_optimize = FunctionToOptimize(
        file_path=main_file,
        function_name="compute",
        parents=(
            FunctionParent(name="OuterClass", type="ClassDef"),
            FunctionParent(name="InnerProcessor", type="ClassDef"),
        ),
    )

    # Test detection directly (context extraction might not work for nested classes)
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        # Create a minimal context for testing
        type(
            "MockContext",
            (),
            {
                "helper_functions": [
                    type(
                        "MockHelper",
                        (),
                        {
                            "qualified_name": "global_helper_1",
                            "only_function_name": "global_helper_1",
                            "fully_qualified_name": "main.global_helper_1",
                            "file_path": main_file,
                            "definition_type": "function",
                        },
                    )(),
                    type(
                        "MockHelper",
                        (),
                        {
                            "qualified_name": "global_helper_2",
                            "only_function_name": "global_helper_2",
                            "fully_qualified_name": "main.global_helper_2",
                            "file_path": main_file,
                            "definition_type": "function",
                        },
                    )(),
                ]
            },
        )(),
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect global_helper_2 as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    expected_unused = {"global_helper_2"}

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )

    # For nested class tests, we'll skip the complete workflow test since nested classes
    # may not be fully supported by the optimizer, but we've verified detection works

    # Also test the complete replace workflow
    # Since this test uses nested classes which might not be fully supported,
    # we'll only test with the mock context for detection but skip the full workflow test
    # The other tests cover the complete workflow comprehensively


def test_multi_file_import_styles(tmp_path: Path) -> None:
    """Test detection with different import styles in multi-file projects."""
    temp_dir = tmp_path

    # Main file
    main_file = temp_dir / "main.py"
    main_file.write_text("""
import utils
from math_helpers import add, multiply
from processors import process_data as pd

def entrypoint_function(n):
    \"\"\"Function using different import styles.\"\"\"
    result1 = utils.compute(n)  # Module.function style
    result2 = add(n, 5)  # Direct import style
    result3 = multiply(n, 2)  # Direct import style
    result4 = pd(n)  # Aliased import style
    return result1 + result2 + result3 + result4
""")

    # Utils file
    utils_file = temp_dir / "utils.py"
    utils_file.write_text("""
def compute(x):
    \"\"\"Utility compute function.\"\"\"
    return x * 10

def unused_util(x):
    \"\"\"This utility function should be unused.\"\"\"
    return x + 100
""")

    # Math helpers file
    math_file = temp_dir / "math_helpers.py"
    math_file.write_text("""
def add(x, y):
    \"\"\"Add two numbers.\"\"\"
    return x + y

def multiply(x, y):
    \"\"\"Multiply two numbers.\"\"\"
    return x * y

def subtract(x, y):
    \"\"\"Subtract function - should be unused.\"\"\"
    return x - y
""")

    # Processors file
    processors_file = temp_dir / "processors.py"
    processors_file.write_text("""
def process_data(x):
    \"\"\"Process data.\"\"\"
    return x ** 2

def clean_data(x):
    \"\"\"Clean data - should be unused.\"\"\"
    return x
""")

    # Optimized version that only uses some functions
    optimized_code = """
```python:main.py
import utils
from math_helpers import add

def entrypoint_function(n):
    \"\"\"Optimized function using fewer helpers.\"\"\"
    result1 = utils.compute(n)  # Still using utils.compute
    result2 = add(n, 5)  # Still using add
    # Inlined multiply: result3 = n * 2
    # Inlined process_data: result4 = n ** 2
    return result1 + result2 + (n * 2) + (n ** 2)
```
"""

    # Create FunctionToOptimize instance
    function_to_optimize = FunctionToOptimize(
        file_path=main_file, function_name="entrypoint_function", parents=()
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect multiply, process_data as unused (at minimum)
    unused_names = {uh.qualified_name for uh in unused_helpers}

    # The exact unused functions may vary based on what helpers are discovered by Jedi
    # At minimum, we expect multiply to be detected as unused since it's not imported
    assert "multiply" in unused_names, (
        "Expected multiply to be detected as unused"
    )
    assert "process_data" in unused_names, (
        "Expected process_data to be detected as unused"
    )
    assert "subtract" not in unused_names, (
        "Expected subtract not to be detected as unused"
    )

    # Also test the complete replace workflow
    # First modify some helper files to simulate optimization changes
    math_file.write_text("""
def add(x, y):
    \"\"\"Add two numbers.\"\"\"
    return x + y

def multiply(x, y):
    \"\"\"Multiply two numbers - MODIFIED VERSION.\"\"\"
    return x * y * 2  # This should be reverted to x * y

def subtract(x, y):
    \"\"\"Subtract function - should be unused.\"\"\"
    return x - y
""")

    # Store original helper code
    original_helper_code = {
        main_file: """
import utils
from math_helpers import add, multiply
from processors import process_data as pd

def entrypoint_function(n):
    \"\"\"Function using different import styles.\"\"\"
    result1 = utils.compute(n)  # Module.function style
    result2 = add(n, 5)  # Direct import style
    result3 = multiply(n, 2)  # Direct import style
    result4 = pd(n)  # Aliased import style
    return result1 + result2 + result3 + result4
""",
        utils_file: utils_file.read_text(),
        math_file: """
def add(x, y):
    \"\"\"Add two numbers.\"\"\"
    return x + y

def multiply(x, y):
    \"\"\"Multiply two numbers.\"\"\"
    return x * y

def subtract(x, y):
    \"\"\"Subtract function - should be unused.\"\"\"
    return x - y
""",
        processors_file: processors_file.read_text(),
    }

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )

    # Check main file content
    main_content = main_file.read_text()
    assert "(n * 2) + (n ** 2)" in main_content, (
        "Entrypoint function should be optimized with inlined calculations"
    )
    assert "from math_helpers import add" in main_content, (
        "Imports should be updated to only include used functions"
    )

    # Verify that unused helper files are reverted if they contained unused functions that were modified
    math_content = math_file.read_text()
    assert "def add(x, y):" in math_content, "add function should still exist"
    # If multiply was unused and modified, it should be reverted
    if "multiply" in unused_names:
        assert "return x * y" in math_content, (
            "multiply should be reverted to original if it was unused"
        )
        assert "return x * y * 2" not in math_content, (
            "multiply should NOT contain the modified version if it was unused"
        )


def test_module_dot_function_import_style(tmp_path: Path) -> None:
    """Test detection when helpers are called via module.function style."""
    temp_dir = tmp_path

    # Main file
    main_file = temp_dir / "main.py"
    main_file.write_text("""
import calculator

def entrypoint_function(n):
    \"\"\"Function using module.function import style.\"\"\"
    result1 = calculator.add_numbers(n, 10)
    result2 = calculator.multiply_numbers(n, 5)
    return result1 + result2
""")

    # Calculator file
    calc_file = temp_dir / "calculator.py"
    calc_file.write_text("""
def add_numbers(x, y):
    \"\"\"Add two numbers.\"\"\"
    return x + y

def multiply_numbers(x, y):
    \"\"\"Multiply two numbers.\"\"\"
    return x * y

def divide_numbers(x, y):
    \"\"\"Divide function - should be unused.\"\"\"
    return x / y
""")

    # Optimized version that only uses add_numbers
    optimized_code = """
```python:main.py
import calculator

def entrypoint_function(n):
    \"\"\"Optimized function that inlines multiply.\"\"\"
    result1 = calculator.add_numbers(n, 10)
    # Inlined: result2 = n * 5
    return result1 + (n * 5)
```
"""

    # Create FunctionToOptimize instance
    function_to_optimize = FunctionToOptimize(
        file_path=main_file, function_name="entrypoint_function", parents=()
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect multiply_numbers and divide_numbers as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}

    # Check that multiply_numbers is detected as unused
    assert "multiply_numbers" in unused_names, (
        f"Expected 'multiply_numbers' to be unused, got: {unused_names}"
    )

    # Also test the complete replace workflow
    # First modify the calculator file to simulate optimization changes
    calc_file.write_text("""
def add_numbers(x, y):
    \"\"\"Add two numbers.\"\"\"
    return x + y

def multiply_numbers(x, y):
    \"\"\"Multiply two numbers - MODIFIED VERSION.\"\"\"
    return x * y * 5  # This should be reverted to x * y

def divide_numbers(x, y):
    \"\"\"Divide function - should be unused.\"\"\"
    return x / y
""")

    # Store original helper code
    original_helper_code = {
        main_file: """
import calculator

def entrypoint_function(n):
    \"\"\"Function using module.function import style.\"\"\"
    result1 = calculator.add_numbers(n, 10)
    result2 = calculator.multiply_numbers(n, 5)
    return result1 + result2
""",
        calc_file: """
def add_numbers(x, y):
    \"\"\"Add two numbers.\"\"\"
    return x + y

def multiply_numbers(x, y):
    \"\"\"Multiply two numbers.\"\"\"
    return x * y

def divide_numbers(x, y):
    \"\"\"Divide function - should be unused.\"\"\"
    return x / y
""",
    }

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )

    # Check main file content
    main_content = main_file.read_text()
    assert "+ (n * 5)" in main_content, (
        "Entrypoint function should be optimized with inlined multiplication"
    )
    assert "import calculator" in main_content, (
        "Calculator import should remain"
    )

    # Check calculator file content - unused functions should be reverted if modified
    calc_content = calc_file.read_text()
    assert "def add_numbers(x, y):" in calc_content, (
        "add_numbers should still exist"
    )
    assert "def multiply_numbers(x, y):" in calc_content, (
        "multiply_numbers should exist"
    )
    assert "def divide_numbers(x, y):" in calc_content, (
        "divide_numbers should remain as original"
    )
    # multiply_numbers should be reverted to original since it's unused
    assert "return x * y" in calc_content, (
        "multiply_numbers should be reverted to original"
    )
    assert "return x * y * 5" not in calc_content, (
        "multiply_numbers should NOT contain the modified version"
    )

    # Also test the complete replace workflow
    # First modify the calculator file to simulate optimization changes
    calc_file.write_text("""
def add_numbers(x, y):
    \"\"\"Add two numbers.\"\"\"
    return x + y

def multiply_numbers(x, y):
    \"\"\"Multiply two numbers - MODIFIED VERSION.\"\"\"
    return x * y * 3  # This should be reverted to x * y

def divide_numbers(x, y):
    \"\"\"Divide function - should be unused.\"\"\"
    return x / y
""")

    # Store original helper code
    original_helper_code = {
        main_file: """
import calculator

def entrypoint_function(n):
    \"\"\"Function using module.function import style.\"\"\"
    result1 = calculator.add_numbers(n, 10)
    result2 = calculator.multiply_numbers(n, 5)
    return result1 + result2
""",
        calc_file: """
def add_numbers(x, y):
    \"\"\"Add two numbers.\"\"\"
    return x + y

def multiply_numbers(x, y):
    \"\"\"Multiply two numbers.\"\"\"
    return x * y

def divide_numbers(x, y):
    \"\"\"Divide function - should be unused.\"\"\"
    return x / y
""",
    }

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )

    # Check main file content
    main_content = main_file.read_text()
    assert "+ (n * 5)" in main_content, (
        "Entrypoint function should be optimized with inlined multiplication"
    )
    assert "import calculator" in main_content, (
        "Calculator import should remain"
    )

    # Check calculator file content - unused functions should be reverted if modified
    calc_content = calc_file.read_text()
    assert "def add_numbers(x, y):" in calc_content, (
        "add_numbers should still exist"
    )
    assert "def multiply_numbers(x, y):" in calc_content, (
        "multiply_numbers should exist"
    )
    assert "def divide_numbers(x, y):" in calc_content, (
        "divide_numbers should remain as original"
    )
    # multiply_numbers should be reverted to original since it's unused
    assert "return x * y" in calc_content, (
        "multiply_numbers should be reverted to original"
    )
    assert "return x * y * 3" not in calc_content, (
        "multiply_numbers should NOT contain the modified version"
    )


def test_static_method_and_class_method(tmp_path: Path) -> None:
    """Test optimization of static methods and class methods."""
    temp_dir = tmp_path

    # Main file with static and class methods
    main_file = temp_dir / "main.py"
    main_file.write_text("""
def utility_function_1(x):
    return x * 2

def utility_function_2(x):
    return x * 3

class MathUtils:
    @staticmethod
    def calculate_static(n):
        \"\"\"Static method that calls utility functions.\"\"\"
        result1 = utility_function_1(n)
        result2 = utility_function_2(n)
        return result1 + result2

    @classmethod
    def calculate_class(cls, n):
        \"\"\"Class method that calls utility functions.\"\"\"
        result1 = utility_function_1(n)
        result2 = utility_function_2(n)
        return result1 - result2
""")

    # Optimized static method that inlines one utility
    optimized_static_code = """
```python:main.py
def utility_function_1(x):
    return x * 2

def utility_function_2(x):
    return x * 3

class MathUtils:
    @staticmethod
    def calculate_static(n):
        \"\"\"Optimized static method.\"\"\"
        result1 = utility_function_1(n)
        return result1 + n * 3  # Inlined utility_function_2

    @classmethod
    def calculate_class(cls, n):
        \"\"\"Class method that calls utility functions.\"\"\"
        result1 = utility_function_1(n)
        result2 = utility_function_2(n)
        return result1 - result2
```
"""

    # Test static method optimization
    function_to_optimize = FunctionToOptimize(
        file_path=main_file,
        function_name="calculate_static",
        parents=(FunctionParent(name="MathUtils", type="ClassDef"),),
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection for static method
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_static_code),
    )

    # Should detect utility_function_2 as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    # The new context pipeline includes the entrypoint itself in helper_functions,
    # so it also appears as "unused" (it doesn't call itself).
    expected_unused = {"utility_function_2", "MathUtils.calculate_static"}

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )

    # Also test the complete replace workflow
    # Update optimized code to include a MODIFIED unused helper
    optimized_static_code_with_modified_helper = """
```python:main.py
def utility_function_1(x):
    return x * 2

def utility_function_2(x):
    return x * 6  # MODIFIED VERSION - should be reverted to x * 3

class MathUtils:
    @staticmethod
    def calculate_static(n):
        \"\"\"Optimized static method.\"\"\"
        result1 = utility_function_1(n)
        return result1 + n * 3  # Inlined utility_function_2

    @classmethod
    def calculate_class(cls, n):
        \"\"\"Class method that calls utility functions.\"\"\"
        result1 = utility_function_1(n)
        result2 = utility_function_2(n)
        return result1 - result2
```
"""

    original_helper_code = {main_file: main_file.read_text()}

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_static_code_with_modified_helper,
        original_helper_code,
        temp_dir,
    )

    # Check final file content
    final_content = main_file.read_text()

    # The static method should be optimized
    assert "result1 + n * 3" in final_content, (
        "Static method should be optimized"
    )

    # utility_function_2 should be reverted to original (x * 3, NOT the modified x * 6)
    assert "return x * 3" in final_content, (
        "utility_function_2 should be reverted to original"
    )
    assert "return x * 6" not in final_content, (
        "utility_function_2 should NOT contain the modified version"
    )

    # utility_function_1 should remain (it's still called)
    assert "def utility_function_1(x):" in final_content, (
        "utility_function_1 should still exist"
    )


def test_async_entrypoint_with_async_helpers(tmp_path: Path) -> None:
    """Test that unused async helper functions are correctly detected when entrypoint is async."""
    temp_dir = tmp_path

    # Main file with async entrypoint and async helpers
    main_file = temp_dir / "main.py"
    main_file.write_text("""
async def async_helper_1(x):
    \"\"\"First async helper function.\"\"\"
    return x * 2

async def async_helper_2(x):
    \"\"\"Second async helper function.\"\"\"
    return x * 3

async def async_entrypoint(n):
    \"\"\"Async entrypoint function that calls async helpers.\"\"\"
    result1 = await async_helper_1(n)
    result2 = await async_helper_2(n)
    return result1 + result2
""")

    # Optimized version that only calls one async helper
    optimized_code = """
```python:main.py
async def async_helper_1(x):
    \"\"\"First async helper function.\"\"\"
    return x * 2

async def async_helper_2(x):
    \"\"\"Second async helper function - should be unused.\"\"\"
    return x * 3

async def async_entrypoint(n):
    \"\"\"Optimized async entrypoint that only calls one helper.\"\"\"
    result1 = await async_helper_1(n)
    return result1 + n * 3  # Inlined async_helper_2
```
"""

    # Create FunctionToOptimize instance for async function
    function_to_optimize = FunctionToOptimize(
        file_path=main_file,
        function_name="async_entrypoint",
        parents=(),
        is_async=True,
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect async_helper_2 as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    expected_unused = {"async_helper_2"}

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )


def test_sync_entrypoint_with_async_helpers(tmp_path: Path) -> None:
    """Test that unused async helper functions are detected when entrypoint is sync."""
    temp_dir = tmp_path

    # Main file with sync entrypoint and async helpers
    main_file = temp_dir / "main.py"
    main_file.write_text("""
import asyncio

async def async_helper_1(x):
    \"\"\"First async helper function.\"\"\"
    return x * 2

async def async_helper_2(x):
    \"\"\"Second async helper function.\"\"\"
    return x * 3

def sync_entrypoint(n):
    \"\"\"Sync entrypoint function that calls async helpers.\"\"\"
    result1 = asyncio.run(async_helper_1(n))
    result2 = asyncio.run(async_helper_2(n))
    return result1 + result2
""")

    # Optimized version that only calls one async helper
    optimized_code = """
```python:main.py
import asyncio

async def async_helper_1(x):
    \"\"\"First async helper function.\"\"\"
    return x * 2

async def async_helper_2(x):
    \"\"\"Second async helper function - should be unused.\"\"\"
    return x * 3

def sync_entrypoint(n):
    \"\"\"Optimized sync entrypoint that only calls one async helper.\"\"\"
    result1 = asyncio.run(async_helper_1(n))
    return result1 + n * 3  # Inlined async_helper_2
```
"""

    # Create FunctionToOptimize instance for sync function
    function_to_optimize = FunctionToOptimize(
        file_path=main_file, function_name="sync_entrypoint", parents=()
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect async_helper_2 as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    expected_unused = {"async_helper_2"}

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )


def test_mixed_sync_and_async_helpers(tmp_path: Path) -> None:
    """Test detection when both sync and async helpers are mixed."""
    temp_dir = tmp_path

    # Main file with mixed sync and async helpers
    main_file = temp_dir / "main.py"
    main_file.write_text("""
import asyncio

def sync_helper_1(x):
    \"\"\"Sync helper function.\"\"\"
    return x * 2

async def async_helper_1(x):
    \"\"\"Async helper function.\"\"\"
    return x * 3

def sync_helper_2(x):
    \"\"\"Another sync helper function.\"\"\"
    return x * 4

async def async_helper_2(x):
    \"\"\"Another async helper function.\"\"\"
    return x * 5

async def mixed_entrypoint(n):
    \"\"\"Async entrypoint function that calls both sync and async helpers.\"\"\"
    sync_result = sync_helper_1(n)
    async_result = await async_helper_1(n)
    sync_result2 = sync_helper_2(n)
    async_result2 = await async_helper_2(n)
    return sync_result + async_result + sync_result2 + async_result2
""")

    # Optimized version that only calls some helpers
    optimized_code = """
```python:main.py
import asyncio

def sync_helper_1(x):
    \"\"\"Sync helper function.\"\"\"
    return x * 2

async def async_helper_1(x):
    \"\"\"Async helper function.\"\"\"
    return x * 3

def sync_helper_2(x):
    \"\"\"Another sync helper function - should be unused.\"\"\"
    return x * 4

async def async_helper_2(x):
    \"\"\"Another async helper function - should be unused.\"\"\"
    return x * 5

async def mixed_entrypoint(n):
    \"\"\"Optimized async entrypoint that only calls some helpers.\"\"\"
    sync_result = sync_helper_1(n)
    async_result = await async_helper_1(n)
    return sync_result + async_result + n * 4 + n * 5  # Inlined both helper_2 functions
```
"""

    # Create FunctionToOptimize instance for async function
    function_to_optimize = FunctionToOptimize(
        file_path=main_file,
        function_name="mixed_entrypoint",
        parents=(),
        is_async=True,
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect both sync_helper_2 and async_helper_2 as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    expected_unused = {"sync_helper_2", "async_helper_2"}

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )


def test_async_class_methods(tmp_path: Path) -> None:
    """Test unused async method detection in classes."""
    temp_dir = tmp_path

    # Main file with class containing async methods
    main_file = temp_dir / "main.py"
    main_file.write_text("""
class AsyncProcessor:
    async def entrypoint_method(self, n):
        \"\"\"Async main method that calls async helper methods.\"\"\"
        result1 = await self.async_helper_method_1(n)
        result2 = await self.async_helper_method_2(n)
        return result1 + result2

    async def async_helper_method_1(self, x):
        \"\"\"First async helper method.\"\"\"
        return x * 2

    async def async_helper_method_2(self, x):
        \"\"\"Second async helper method.\"\"\"
        return x * 3

    def sync_helper_method(self, x):
        \"\"\"Sync helper method.\"\"\"
        return x * 4
""")

    # Optimized version that only calls one async helper
    optimized_code = """
```python:main.py
class AsyncProcessor:
    async def entrypoint_method(self, n):
        \"\"\"Optimized async method that only calls one helper.\"\"\"
        result1 = await self.async_helper_method_1(n)
        return result1 + n * 3  # Inlined async_helper_method_2

    async def async_helper_method_1(self, x):
        \"\"\"First async helper method.\"\"\"
        return x * 2

    async def async_helper_method_2(self, x):
        \"\"\"Second async helper method - should be unused.\"\"\"
        return x * 3

    def sync_helper_method(self, x):
        \"\"\"Sync helper method - should be unused.\"\"\"
        return x * 4
```
"""

    # Create FunctionToOptimize instance for async class method
    function_to_optimize = FunctionToOptimize(
        file_path=main_file,
        function_name="entrypoint_method",
        parents=(FunctionParent(name="AsyncProcessor", type="ClassDef"),),
        is_async=True,
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect async_helper_method_2 as unused (sync_helper_method may not be discovered as helper)
    unused_names = {uh.qualified_name for uh in unused_helpers}
    # The new context pipeline includes the entrypoint itself in helper_functions,
    # so it also appears as "unused" (it doesn't call itself).
    expected_unused = {
        "AsyncProcessor.async_helper_method_2",
        "AsyncProcessor.entrypoint_method",
    }

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )


def test_async_helper_revert_functionality(tmp_path: Path) -> None:
    """Test that unused async helper functions are correctly reverted to original definitions."""
    temp_dir = tmp_path

    # Main file with async functions
    main_file = temp_dir / "main.py"
    main_file.write_text("""
async def async_helper_1(x):
    \"\"\"First async helper function.\"\"\"
    return x * 2

async def async_helper_2(x):
    \"\"\"Second async helper function.\"\"\"
    return x * 3

async def async_entrypoint(n):
    \"\"\"Async entrypoint function that calls async helpers.\"\"\"
    result1 = await async_helper_1(n)
    result2 = await async_helper_2(n)
    return result1 + result2
""")

    # Optimized version that only calls one helper and modifies the unused one
    optimized_code = """
```python:main.py
async def async_helper_1(x):
    \"\"\"First async helper function.\"\"\"
    return x * 2

async def async_helper_2(x):
    \"\"\"Modified async helper function - should be reverted.\"\"\"
    return x * 10  # This change should be reverted

async def async_entrypoint(n):
    \"\"\"Optimized async entrypoint that only calls one helper.\"\"\"
    result1 = await async_helper_1(n)
    return result1 + n * 3  # Inlined async_helper_2
```
"""

    # Create FunctionToOptimize instance for async function
    function_to_optimize = FunctionToOptimize(
        file_path=main_file,
        function_name="async_entrypoint",
        parents=(),
        is_async=True,
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Store original helper code
    original_helper_code = {main_file: main_file.read_text()}

    # Apply optimization and test reversion
    apply_optimized_code(
        function_to_optimize,
        code_context,
        optimized_code,
        original_helper_code,
        temp_dir,
    )

    # Check final file content
    final_content = main_file.read_text()

    # The entrypoint should be optimized
    assert "result1 + n * 3" in final_content, (
        "Async entrypoint function should be optimized"
    )

    # async_helper_2 should be reverted to original (return x * 3, not x * 10)
    assert "return x * 3" in final_content, (
        "async_helper_2 should be reverted to original"
    )
    assert "return x * 10" not in final_content, (
        "async_helper_2 should not contain the modified version"
    )

    # async_helper_1 should remain (it's still called)
    assert "async def async_helper_1(x):" in final_content, (
        "async_helper_1 should still exist"
    )


def test_recursive_helper_function_not_detected_as_unused(
    tmp_path: Path,
) -> None:
    """Test that recursive helper functions are NOT incorrectly detected as unused."""
    temp_dir = tmp_path

    # Main file with recursive helper function
    main_file = temp_dir / "main.py"
    main_file.write_text("""
def gcd_recursive(a: int, b: int) -> int:
    \"\"\"Calculate greatest common divisor using Euclidean algorithm with recursion.\"\"\"
    if b == 0:
        return a
    return gcd_recursive(b, a % b)
""")

    # Optimized version that still uses the recursive helper
    optimized_code = """
```python:main.py
def gcd_recursive(a: int, b: int) -> int:
    \"\"\"Calculate greatest common divisor using Euclidean algorithm with recursion.\"\"\"
    if b == 0:
        return a
    return gcd_recursive(b, a % b)
```
"""

    # Create FunctionToOptimize instance
    function_to_optimize = FunctionToOptimize(
        file_path=main_file, function_name="gcd_recursive", parents=()
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should NOT detect gcd_recursive as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}

    assert "gcd_recursive" not in unused_names, (
        f"Recursive function gcd_recursive should NOT be detected as unused, but got unused: {unused_names}"
    )


def test_async_generators_and_coroutines(tmp_path: Path) -> None:
    """Test detection with async generators and coroutines."""
    temp_dir = tmp_path

    # Main file with async generators and coroutines
    main_file = temp_dir / "main.py"
    main_file.write_text("""
import asyncio

async def async_generator_helper(n):
    \"\"\"Async generator helper.\"\"\"
    for i in range(n):
        yield i * 2

async def coroutine_helper(x):
    \"\"\"Coroutine helper.\"\"\"
    await asyncio.sleep(0.1)
    return x * 3

async def another_coroutine_helper(x):
    \"\"\"Another coroutine helper.\"\"\"
    await asyncio.sleep(0.1)
    return x * 4

async def async_entrypoint_with_generators(n):
    \"\"\"Async entrypoint function that uses generators and coroutines.\"\"\"
    results = []
    async for value in async_generator_helper(n):
        results.append(value)

    final_result = await coroutine_helper(sum(results))
    another_result = await another_coroutine_helper(n)
    return final_result + another_result
""")

    # Optimized version that doesn't use one of the coroutines
    optimized_code = """
```python:main.py
import asyncio

async def async_generator_helper(n):
    \"\"\"Async generator helper.\"\"\"
    for i in range(n):
        yield i * 2

async def coroutine_helper(x):
    \"\"\"Coroutine helper.\"\"\"
    await asyncio.sleep(0.1)
    return x * 3

async def another_coroutine_helper(x):
    \"\"\"Another coroutine helper - should be unused.\"\"\"
    await asyncio.sleep(0.1)
    return x * 4

async def async_entrypoint_with_generators(n):
    \"\"\"Optimized async entrypoint that inlines one coroutine.\"\"\"
    results = []
    async for value in async_generator_helper(n):
        results.append(value)

    final_result = await coroutine_helper(sum(results))
    return final_result + n * 4  # Inlined another_coroutine_helper
```
"""

    # Create FunctionToOptimize instance for async function
    function_to_optimize = FunctionToOptimize(
        file_path=main_file,
        function_name="async_entrypoint_with_generators",
        parents=(),
        is_async=True,
    )

    # Get original code context
    code_context = get_code_optimization_context(
        function_to_optimize, temp_dir
    )

    # Test unused helper detection
    unused_helpers = detect_unused_helper_functions(
        function_to_optimize,
        code_context,
        CodeStringsMarkdown.parse_markdown_code(optimized_code),
    )

    # Should detect another_coroutine_helper as unused
    unused_names = {uh.qualified_name for uh in unused_helpers}
    expected_unused = {"another_coroutine_helper"}

    assert unused_names == expected_unused, (
        f"Expected unused: {expected_unused}, got: {unused_names}"
    )
