"""Tests for comment cleaner post-processing step.

The comment cleaner removes or reverts comments that are not immediately before,
within, or after modified code lines. This prevents LLMs from adding unnecessary
documentation to unchanged code.
"""

import ast

import libcst as cst

from core.languages.python.optimizer.postprocess import clean_extraneous_comments


def assert_code_unchanged(optimized: str, result: str) -> None:
    """Verify that the actual code (not comments) is identical between optimized and result.

    This ensures the comment cleaner only modifies comments and doesn't accidentally
    change the actual code logic from what the optimizer produced.
    """
    # Parse both as AST and compare the unparsed versions
    # This strips all comments and formatting, leaving only the semantic code
    optimized_ast = ast.parse(optimized)
    result_ast = ast.parse(result)

    # Unparse to get normalized code without comments
    optimized_code = ast.unparse(optimized_ast)
    result_code = ast.unparse(result_ast)

    assert optimized_code == result_code, (
        f"Code was modified (not just comments)!\nOptimized code: {optimized_code}\nResult code: {result_code}"
    )


def test_keep_comments_near_modified_code():
    """Test Case 1: Keep comments that are within modified code region."""
    original = """def calculate(x, y):
    return x + y"""

    optimized = """def calculate(x, y):
    # Using bit shift for faster multiplication by 2
    return (x + y) << 1"""

    expected = """def calculate(x, y):
    # Using bit shift for faster multiplication by 2
    return (x + y) << 1"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_remove_comments_added_to_unchanged_code():
    """Test Case 2: Remove comments added to unchanged code lines."""
    original = """def calculate(x, y):
    z = x + y
    return z * 2"""

    optimized = """def calculate(x, y):
    # Store intermediate result
    z = x + y  # This computes the sum
    return z << 1"""

    expected = """def calculate(x, y):
    z = x + y
    return z << 1"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_keep_comments_immediately_before_modified_code():
    """Test Case 3: Remove comments before the block of the modified code."""
    original = """def process(data):
    result = []
    for item in data:
        result.append(item * 2)
    return result"""

    optimized = """def process(data):
    result = []
    # Using list comprehension for better performance
    for item in data:
        result.extend([item * 2])
    return result"""

    expected = """def process(data):
    result = []
    for item in data:
        result.extend([item * 2])
    return result"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_keep_comments_immediately_after_modified_code():
    """Test Case 4: Keep inline comments on modified lines."""
    original = """def validate(x):
    if x > 0:
        return True
    return False"""

    optimized = """def validate(x):
    if x > 0:
        return False  # Optimization applied here
    return False"""

    expected = """def validate(x):
    if x > 0:
        return False  # Optimization applied here
    return False"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_remove_comments_far_from_modifications():
    """Test Case 5: Remove comments that are far from any modifications."""
    original = """def complex_calc(a, b, c):
    x = a + b
    y = x * c
    return y"""

    optimized = """def complex_calc(a, b, c):
    # Calculate intermediate sum
    x = a + b
    # Perform final multiplication using bit shift
    y = x << c
    # Return the result
    return y"""

    expected = """def complex_calc(a, b, c):
    x = a + b
    # Perform final multiplication using bit shift
    y = x << c
    # Return the result
    return y"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_use_modified_existing_comments_next_to_changes():
    original = """def helper(n):
    # Original comment about n
    result = n * 2
    return result"""

    optimized = """def helper(n):
    # LLM modified this comment to explain n differently
    result = n << 1
    return result"""

    expected = """def helper(n):
    # LLM modified this comment to explain n differently
    result = n << 1
    return result"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_multiple_modifications_with_mixed_comments():
    """Test Case 7: Handle multiple modifications with mixed comment scenarios."""
    original = """def multi_op(a, b, c):
    x = a + b
    y = x * 2
    z = y + c
    return z"""

    optimized = """def multi_op(a, b, c):
    # Sum a and b
    x = a + b
    # Double using bit shift
    y = x << 1
    # Add c
    z = y + c
    return z"""

    expected = """def multi_op(a, b, c):
    x = a + b
    # Double using bit shift
    y = x << 1
    # Add c
    z = y + c
    return z"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_no_modifications_remove_all_new_comments():
    """Test Case 9: Remove all new comments when no code was actually modified."""
    original = """def unchanged(x):
    return x * 2"""

    optimized = """def unchanged(x):
    # This computes double
    return x * 2"""

    expected = """def unchanged(x):
    return x * 2"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_block_comment_preservation_near_modifications():
    """Test Case 10: Handle multi-line comments near modifications."""
    original = """def block_test(n):
    if n > 0:
        return n * 2
    return 0"""

    optimized = """def block_test(n):
    # This function handles positive numbers
    # by doubling them efficiently
    if n > 0:
        # Using bit shift optimization
        return n << 1
    return 0"""

    expected = """def block_test(n):
    if n > 0:
        # Using bit shift optimization
        return n << 1
    return 0"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_preserve_original_comments_on_modified_lines():
    """Test that original comments on modified lines are preserved."""
    original = """def func(x):
    # Original comment
    result = x * 2  # inline original
    return result"""

    optimized = """def func(x):
    # Modified comment
    result = x << 1  # inline modified
    return result"""

    expected = """def func(x):
    # Modified comment
    result = x << 1  # inline modified
    return result"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_empty_file():
    """Test handling of empty files."""
    original = ""
    optimized = ""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    # Empty file case - both original and result are empty, so AST check would pass trivially
    assert result.code == ""


def test_nested_function_with_comments():
    """Test handling of comments in nested functions."""
    original = """def outer():
    def inner(x):
        return x * 2
    return inner"""

    optimized = """def outer():
    # Nested function for doubling
    def inner(x):
        # Using bit shift
        return x << 1
    return inner"""

    expected = """def outer():
    def inner(x):
        # Using bit shift
        return x << 1
    return inner"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_class_with_method_modifications():
    """Test handling of comments in class methods."""
    original = """class Calculator:
    def calculate(self, x):
        return x * 2"""

    optimized = """class Calculator:
    # Main calculator class
    def calculate(self, x):
        # Optimized calculation
        return x << 1"""

    expected = """class Calculator:
    def calculate(self, x):
        # Optimized calculation
        return x << 1"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_strings_with_hash_symbols_not_treated_as_comments():
    """Test that hash symbols and comment-like text inside strings are preserved as code."""
    original = """def process_text(data):
    message = "Processing # of items"
    multiline = \"\"\"
    This is a multi-line string
    # This looks like a comment but it's in a string
    \"\"\"
    result = data.split("#")
    return result"""

    optimized = """def process_text(data):
    # Add helpful message
    message = "Processing # of items"
    # Multi-line string with comment-like text
    multiline = \"\"\"
    This is a multi-line string
    # This looks like a comment but it's in a string
    \"\"\"
    # Split by hash symbol
    result = data.split("#")
    return result"""

    expected = """def process_text(data):
    message = "Processing # of items"
    multiline = \"\"\"
    This is a multi-line string
    # This looks like a comment but it's in a string
    \"\"\"
    result = data.split("#")
    return result"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_strings_with_hash_before_code_changes():
    """Test that strings with # immediately before code changes are preserved correctly."""
    original = """def format_and_process(data):
    header = "# Header information"
    info = \"\"\"
    Documentation:
    # This is documentation, not a comment
    # Another line in docs
    \"\"\"
    result = data * 2
    return result"""

    optimized = """def format_and_process(data):
    # String with hash symbol
    header = "# Header information"
    # Multi-line documentation string
    info = \"\"\"
    Documentation:
    # This is documentation, not a comment
    # Another line in docs
    \"\"\"
    # Optimized to use bit shift
    result = data << 1
    return result"""

    expected = """def format_and_process(data):
    header = "# Header information"
    info = \"\"\"
    Documentation:
    # This is documentation, not a comment
    # Another line in docs
    \"\"\"
    # Optimized to use bit shift
    result = data << 1
    return result"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_preserve_removed_original_comments_far_from_changes():
    """Test that original comments removed by LLM are restored when immediately before unchanged code."""
    original = """def calculate_stats(numbers):
    # This is important context about the function
    total = 0
    # Track the count
    count = 0

    for num in numbers:
        # Process each number
        total += num
        count += 1

    return total / count"""

    optimized = """def calculate_stats(numbers):
    total = 0
    count = 0

    for num in numbers:
        total += num
        count += 1

    # Optimized: use built-in functions
    return sum(numbers) / len(numbers)"""

    # Expected behavior:
    # - Original comments immediately before unchanged lines (total=0, count=0, etc.) are restored
    # - The blank line between count=0 and for statement: Since both lines around it are unchanged,
    #   and the blank line was in the original, it should be restored
    # - New comment "Optimized: use built-in functions" before the modified return should be kept
    expected = """def calculate_stats(numbers):
    # This is important context about the function
    total = 0
    # Track the count
    count = 0

    for num in numbers:
        # Process each number
        total += num
        count += 1

    # Optimized: use built-in functions
    return sum(numbers) / len(numbers)"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_restore_inline_comments_on_unchanged_lines():
    """Test that inline comments on unchanged lines are preserved from original."""
    original = """def process(value):
    result = value * 2  # Double the input
    output = result + 10  # Add offset
    return output"""

    optimized = """def process(value):
    result = value * 2
    output = result + 10
    # Final step: return the computed value
    return output << 1"""

    expected = """def process(value):
    result = value * 2  # Double the input
    output = result + 10  # Add offset
    # Final step: return the computed value
    return output << 1"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_restore_multiline_comment_block_before_unchanged_code():
    """Test current behavior: only the immediately preceding comment is restored.

    Note: Currently, when multiple consecutive comment lines precede unchanged code,
    only the comment immediately before the code line is restored. This is a known
    limitation - full multi-line comment block restoration would require more
    sophisticated tracking of comment groups.
    """
    original = """def calculate(x, y):
    # This function performs a calculation
    # It takes two parameters and combines them
    # The logic here is intentionally simple
    result = x + y
    return result * 2"""

    optimized = """def calculate(x, y):
    result = x + y
    # Optimized multiplication
    return result << 1"""

    # Current behavior: multi-line comment blocks are not fully restored
    # Only the immediately preceding comment would be restored if the code line was unchanged
    # In this case, result = x + y is unchanged, so comments before it should be restored
    expected = """def calculate(x, y):
    # This function performs a calculation
    # It takes two parameters and combines them
    # The logic here is intentionally simple
    result = x + y
    # Optimized multiplication
    return result << 1"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_mixed_original_and_new_comments():
    """Test handling when LLM both removes original comments and adds new ones.

    The LLM replaces original comments with new ones. The current behavior:
    - If the line is unchanged, remove the new comment (not immediately before a change)
    - Original comments are only restored if immediately before the unchanged line
    """
    original = """def validate_input(data):
    # Check if data is valid
    if not data:
        return False
    # Process the data
    cleaned = data.strip()
    return len(cleaned) > 0"""

    optimized = """def validate_input(data):
    # First validation step
    if not data:
        return False
    # Strip whitespace for better accuracy
    cleaned = data.strip()
    # Use bool conversion for cleaner code
    return bool(cleaned)"""

    # Current behavior: new comments before unchanged lines are removed,
    # but original comments are not fully restored (limitation)
    expected = """def validate_input(data):
    # Check if data is valid
    if not data:
        return False
    # Process the data
    cleaned = data.strip()
    # Use bool conversion for cleaner code
    return bool(cleaned)"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_restore_comments_in_loop_bodies():
    """Test that comments inside loops are preserved for unchanged code."""
    original = """def sum_evens(numbers):
    total = 0
    for num in numbers:
        # Only process even numbers
        if num % 2 == 0:
            total += num
    return total"""

    optimized = """def sum_evens(numbers):
    total = 0
    for num in numbers:
        if num % 2 == 0:
            total += num
    # Optimized return using ternary
    return total if total > 0 else 0"""

    expected = """def sum_evens(numbers):
    total = 0
    for num in numbers:
        # Only process even numbers
        if num % 2 == 0:
            total += num
    # Optimized return using ternary
    return total if total > 0 else 0"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_restore_comments_with_special_characters():
    """Test current behavior with comments containing special characters.

    Similar to multi-line comment blocks, only the immediately preceding comment
    is restored. Multiple consecutive comments are not fully restored.
    """
    original = """def format_string(text):
    # TODO: handle edge cases like empty strings
    # NOTE: this uses str.format() for compatibility
    formatted = "Result: {}".format(text)
    return formatted"""

    optimized = """def format_string(text):
    formatted = "Result: {}".format(text)
    # Enhanced: strip whitespace
    return formatted.strip()"""

    # Current behavior: multi-line comments before unchanged code not fully restored
    expected = """def format_string(text):
    # TODO: handle edge cases like empty strings
    # NOTE: this uses str.format() for compatibility
    formatted = "Result: {}".format(text)
    # Enhanced: strip whitespace
    return formatted.strip()"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_preserve_docstring_and_restore_inline_comments():
    """Test that docstrings are preserved and inline comments are restored."""
    original = """def helper(value):
    \"\"\"Helper function for processing values.\"\"\"
    intermediate = value * 3  # Triple the value
    return intermediate"""

    optimized = """def helper(value):
    \"\"\"Helper function for processing values.\"\"\"
    # Calculate intermediate result
    intermediate = value * 3
    # Return with optimization
    return intermediate + 1"""

    expected = """def helper(value):
    \"\"\"Helper function for processing values.\"\"\"
    intermediate = value * 3  # Triple the value
    # Return with optimization
    return intermediate + 1"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_restore_comments_after_conditional_branches():
    """Test that comments in conditional branches are handled correctly."""
    original = """def check_status(code):
    if code == 200:
        # Success case
        return "OK"
    elif code == 404:
        # Not found case
        return "Not Found"
    else:
        return "Error\""""

    optimized = """def check_status(code):
    if code == 200:
        return "OK"
    elif code == 404:
        return "Not Found"
    else:
        # Default error handling with logging
        return "Unknown Error\""""

    expected = """def check_status(code):
    if code == 200:
        # Success case
        return "OK"
    elif code == 404:
        # Not found case
        return "Not Found"
    else:
        # Default error handling with logging
        return "Unknown Error\""""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()


def test_comprehensive_example_with_multiple_comment_scenarios():
    """Comprehensive test with multiple comment scenarios in realistic code."""
    original = """def process_data(items, threshold=10):
    # Initialize results
    results = []
    total = 0

    for item in items:
        if item > threshold:
            value = item * 2
            total += value
            results.append(value)

    return results, total"""

    optimized = """def process_data(items, threshold=10):
    # Initialize results list
    results = []
    # Track running total
    total = 0

    # Process each item in the list
    for item in items:
        # Check if item exceeds threshold
        if item > threshold:
            # Use bit shift for faster multiplication
            value = item << 1  # Double the value
            # Add to running total
            total += value
            results.append(value)  # Store result

    # Return both results and total
    return results, total"""

    # Expected behavior:
    # - "Initialize results list" - removed (modified comment on unchanged line)
    # - "Track running total" - removed (comment before unchanged line)
    # - "Process each item in the list" - removed (comment before unchanged line)
    # - "Check if item exceeds threshold" - removed (comment before unchanged line)
    # - "Use bit shift for faster multiplication" - KEPT (immediately before modified line)
    # - "Double the value" inline - KEPT (on modified line)
    # - "Add to running total" - KEPT (part of replaced block with the optimization)
    # - "Store result" inline - removed (inline on unchanged line)
    # - "Return both results and total" - removed (comment before unchanged line)
    expected = """def process_data(items, threshold=10):
    # Initialize results
    results = []
    total = 0

    for item in items:
        if item > threshold:
            # Use bit shift for faster multiplication
            value = item << 1  # Double the value
            # Add to running total
            total += value
            results.append(value)

    return results, total"""

    original_module = cst.parse_module(original)
    optimized_module = cst.parse_module(optimized)
    result = clean_extraneous_comments(original_module, optimized_module)

    assert_code_unchanged(optimized, result.code)
    assert result.code.strip() == expected.strip()
