import dataclasses

from aiservice.common.markdown_utils import split_markdown_code
from core.languages.python.optimizer.context_utils.optimizer_context import (
    BaseOptimizerContext,
    MultiOptimizerContext,
    SingleOptimizerContext,
)
from core.languages.python.optimizer.context_utils.refiner_context import BaseRefinerContext, RefinementContextData
from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
from core.shared.context_helpers import group_code

MULTI_CONTEXT_SPLITTER_PREFIX = ""


@dataclasses.dataclass
class OptimizerTestCase:
    llm_response: str
    original_code: str
    expected_optimized_code: str
    expected_explanation: str


@dataclasses.dataclass()
class RefinerTestCase:
    llm_response: str
    optimized_code: str
    expected_refined_code: str


def create_optimizer_context(code: str, diff_method: DiffMethod = DiffMethod.NO_DIFF) -> BaseOptimizerContext:
    return BaseOptimizerContext.get_dynamic_context("", "", code, diff_method)


def create_refiner_context(optimized_code: str) -> BaseRefinerContext:
    return BaseRefinerContext.get_dynamic_context(
        RefinementContextData(
            original_source_code="",
            original_line_profiler_results="",
            original_code_runtime="",
            optimized_source_code=optimized_code,
            read_only_dependency_code="",
            optimized_line_profiler_results="",
            optimized_code_runtime="",
            speedup="",
            optimized_explanation="",
        ),
        "",
        "",
    )


optimizer_with_search_and_replace_test_cases: list[OptimizerTestCase] = [
    OptimizerTestCase(
        original_code="""```python:app/sorter.py
def sorter(arr):
    for i in range(len(arr)):
        for j in range(len(arr) - 1):
            if arr[j] > arr[j + 1]:
                temp = arr[j]
                arr[j] = arr[j + 1]
                arr[j + 1] = temp
    return arr
```
""",
        llm_response="""<replace_in_file>
<path>app/sorter.py</path>
<diff>
<<<<<<< SEARCH
def sorter(arr):
    for i in range(len(arr)):
        for j in range(len(arr) - 1):
            if arr[j] > arr[j + 1]:
                temp = arr[j]
                arr[j] = arr[j + 1]
                arr[j + 1] = temp
    return arr
=======
def sorter(arr):
    n = len(arr)
    for i in range(n):
        # Last i elements are already in place
        for j in range(0, n - i - 1):
            if arr[j] > arr[j + 1]:
                # Swap without an unnecessary temp variable
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
    return arr
>>>>>>> REPLACE
</diff>
</replace_in_file>

Rewrote the double for-loop to avoid redundant comparisons as the largest elements bubble to the end after each pass, reducing iterations. Used tuple assignment for swapping to reduce bytecode instructions and temporary variable usage. Behavior and signature remain unchanged.
""",
        expected_optimized_code="""```python:app/sorter.py
def sorter(arr):
    n = len(arr)
    for i in range(n):
        # Last i elements are already in place
        for j in range(0, n - i - 1):
            if arr[j] > arr[j + 1]:
                # Swap without an unnecessary temp variable
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
    return arr
```""",
        expected_explanation="""Rewrote the double for-loop to avoid redundant comparisons as the largest elements bubble to the end after each pass, reducing iterations. Used tuple assignment for swapping to reduce bytecode instructions and temporary variable usage. Behavior and signature remain unchanged.""",
    )
]


single_optimizer_test_cases: list[OptimizerTestCase] = [
    OptimizerTestCase(
        original_code="",
        llm_response="""Here’s a faster version of your `sorter` function.
Your implementation is a non-optimized bubble sort using O(n²) time always.
For better performance, I’ll use Python’s built-in `.sort()` method, which is highly optimized (Timsort: O(n log n)), but will still print debugging outputs as you request.

```python
def sorter(arr):
    print("codeflash stdout: Sorting list")
    arr.sort()
    print(f"result: {arr}")
    return arr
```
""",
        expected_optimized_code="""def sorter(arr):
    print("codeflash stdout: Sorting list")
    arr.sort()
    print(f"result: {arr}")
    return arr
""",
        expected_explanation="""Here’s a faster version of your `sorter` function.
Your implementation is a non-optimized bubble sort using O(n²) time always.
For better performance, I’ll use Python’s built-in `.sort()` method, which is highly optimized (Timsort: O(n log n)), but will still print debugging outputs as you request.""",
    ),
    # multiple code blocks within same response, should pick first one, update: we are stating in the prompt it should write a short summary of the changes then provide the code so it will be a single block always
    OptimizerTestCase(
        original_code="",
        llm_response="""
- The original code uses a bubble sort (`O(n^2)`). `list.sort()` uses Timsort (`O(n log n)`), which is far faster.
- No additional memory is required (it's in-place).
- Print statements remain as before.
- No change to function signature or return value.

**If you want to preserve the bubble sort structure but still optimize significantly, you can at least avoid repeated `len(arr)` calls and skip unnecessary passes when no swaps are made (standard bubble sort optimization):**

```python
def sorter(arr):
    print("codeflash stdout: Sorting list")
    arr.sort()
    print(f"result: {arr}")
    return arr
```
""",
        expected_optimized_code="""def sorter(arr):
    print("codeflash stdout: Sorting list")
    arr.sort()
    print(f"result: {arr}")
    return arr
""",
        expected_explanation="""- The original code uses a bubble sort (`O(n^2)`). `list.sort()` uses Timsort (`O(n log n)`), which is far faster.
- No additional memory is required (it's in-place).
- Print statements remain as before.
- No change to function signature or return value.

**If you want to preserve the bubble sort structure but still optimize significantly, you can at least avoid repeated `len(arr)` calls and skip unnecessary passes when no swaps are made (standard bubble sort optimization):**""",
    ),
]

single_refiner_test_cases: list[RefinerTestCase] = [
    RefinerTestCase(
        llm_response="""Some explanation by LLM

<replace_in_file>
<path>optimized_source_code</path>
<diff>
<<<<<<< SEARCH
                arr[j], arr[j + 1] = arr[j + 1], arr[j]  # Swap in-place
=======
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
>>>>>>> REPLACE

<<<<<<< SEARCH
            break  # Stop early if sorted
=======
            break
>>>>>>> REPLACE
</diff>
</replace_in_file>""",
        optimized_code="""def sorter(arr):
    print("codeflash stdout: Sorting list")
    n = len(arr)
    for i in range(n):
        swapped = False
        for j in range(0, n - i - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]  # Swap in-place
                swapped = True
        if not swapped:
            break  # Stop early if sorted
    print(f"result: {arr}")
    return arr""",
        expected_refined_code="""def sorter(arr):
    print("codeflash stdout: Sorting list")
    n = len(arr)
    for i in range(n):
        swapped = False
        for j in range(0, n - i - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
                swapped = True
        if not swapped:
            break
    print(f"result: {arr}")
    return arr""",
    )
]

multi_optimizer_test_cases: list[OptimizerTestCase] = [
    OptimizerTestCase(
        # single file with the multi-context
        llm_response="""Here’s a faster, more optimized version, replacing the inefficient bubble sort with the built-in, highly optimized `sort()` method. The function signature and effect are preserved.
This change dramatically improves speed (from O(n^2) to O(n log n)) and reduces memory and CPU usage. Comments are preserved where relevant.
```python
def sorter(arr):
    print("codeflash stdout: Sorting list")
    arr.sort() # much faster than manual bubble sort
    print(f"result: {arr}")
    return arr
```
""",
        original_code="""```python:relative/path/to/bubble_sort.py
def sorter(arr):
    print("codeflash stdout: Sorting list")
    for i in range(len(arr)):
        for j in range(len(arr) - 1):
            if arr[j] > arr[j + 1]:
                temp = arr[j]
                arr[j] = arr[j + 1]
                arr[j + 1] = temp
    print(f"result: {arr}")
    return arr
```
""",
        expected_optimized_code="""```python:relative/path/to/bubble_sort.py
def sorter(arr):
    print("codeflash stdout: Sorting list")
    arr.sort() # much faster than manual bubble sort
    print(f"result: {arr}")
    return arr
```""",
        expected_explanation="""Here’s a faster, more optimized version, replacing the inefficient bubble sort with the built-in, highly optimized `sort()` method. The function signature and effect are preserved.
This change dramatically improves speed (from O(n^2) to O(n log n)) and reduces memory and CPU usage. Comments are preserved where relevant.""",
    ),
    OptimizerTestCase(
        llm_response="""**Explanation of Optimizations:**

For math_operations.py:
- Replaced manual power calculation with built-in `**` operator
- Simplified factorial using math.factorial()
- Added type hints for better code clarity

For string_utils.py:
- Used str.join() instead of += in loops for better performance
- Simplified palindrome check using string slicing
- Added docstrings for better documentation

```python:relative/path/to/math_operations.py
import math

def calculate_power(base: float, exponent: float) -> float:
    return base ** exponent

def calculate_factorial(n: int) -> int:
    return math.factorial(n)
```

```python:relative/path/to/string_utils.py
def concatenate_strings(strings: list[str]) -> str:
    '''Concatenate a list of strings efficiently'''
    return ''.join(strings)

def is_palindrome(s: str) -> bool:
    '''Check if a string is a palindrome'''
    return s == s[::-1]
```
""",
        original_code="""```python:relative/path/to/math_operations.py
def calculate_power(base, exponent):
    result = 1
    for _ in range(exponent):
        result *= base
    return result

def calculate_factorial(n):
    if n == 0:
        return 1
    else:
        return n * calculate_factorial(n - 1)
```
```python:relative/path/to/string_utils.py
def concatenate_strings(strings):
    result = ""
    for s in strings:
        result += s
    return result

def is_palindrome(s):
    for i in range(len(s)//2):
        if s[i] != s[len(s)-1-i]:
            return False
    return True
```
""",
        expected_optimized_code="""```python:relative/path/to/math_operations.py
import math


def calculate_power(base: float, exponent: float) -> float:
    return base ** exponent

def calculate_factorial(n: int) -> int:
    return math.factorial(n)
```
```python:relative/path/to/string_utils.py
def concatenate_strings(strings: list[str]) -> str:
    '''Concatenate a list of strings efficiently'''
    return ''.join(strings)

def is_palindrome(s: str) -> bool:
    '''Check if a string is a palindrome'''
    return s == s[::-1]
```""",
        expected_explanation="""**Explanation of Optimizations:**

For math_operations.py:
- Replaced manual power calculation with built-in `**` operator
- Simplified factorial using math.factorial()
- Added type hints for better code clarity

For string_utils.py:
- Used str.join() instead of += in loops for better performance
- Simplified palindrome check using string slicing
- Added docstrings for better documentation""",
    ),
    OptimizerTestCase(
        llm_response="""### Explanation

**Optimizations Applied:**

- **Inlined swap logic:** Instead of calling the `swap` function inside the innermost loop (which creates substantial function call overhead), I replaced its use with direct tuple assignment (`arr[j], arr[j + 1] = arr[j + 1], arr[j]`). This eliminates unnecessary function calls and stack frames in a hot path, significantly reducing the total runtime.
- **Early termination in Bubble Sort:** Added a `swapped` flag to exit early if no swaps were performed in an inner pass. This prevents unnecessary iterations when the list is already sorted, reducing average case and best case time complexity from always O(n²) to O(n) for already sorted or nearly sorted lists.

**Performance Impact:**
- Reduced hottest-path function call overhead.
- Reduced unnecessary iterations for sorted/permuted lists.
- `swap` in `utils.py` remains for API compatibility but is now unused in the optimized sort.

---

```python:app/utils.py
def swap(val1, val2):
    # create a helper function for this to count down the number of swaps happened
    # (Unchanged: still a tuple swap for API compatibility, no longer used in sorter)
    temp = val1
    val1 = val2
    val2 = temp
    return (val1, val2)
```

```python:app/main.py
from app.utils import swap

def sorter_test_final10(arr):
    print("codeflash stdout: Sorting list")
    n = len(arr)
    for i in range(n):
        swapped = False  # Optimization: Track if any swaps occurred in this pass
        for j in range(n - 1):
            if arr[j] > arr[j + 1]:
                # Fast tuple swap: avoids function call overhead
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
                swapped = True
        if not swapped:
            break  # Stop if the list is already sorted
    print(f"result: {arr}")
    return arr
```""",
        original_code="""```python:app/utils.py
def swap(val1, val2):
    # create a helper function for this to count down the number of swaps happened
    temp = val1
    val1 = val2
    val2 = temp
    return (val1, val2)
```
```python:app/main.py
from app.utils import swap

def sorter_test_final10(arr):
    print("codeflash stdout: Sorting list")
    for i in range(len(arr)):
        for j in range(len(arr) - 1):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = swap(arr[j], arr[j + 1])
    print(f"result: {arr}")
    return arr
```
""",
        expected_explanation="""### Explanation

**Optimizations Applied:**

- **Inlined swap logic:** Instead of calling the `swap` function inside the innermost loop (which creates substantial function call overhead), I replaced its use with direct tuple assignment (`arr[j], arr[j + 1] = arr[j + 1], arr[j]`). This eliminates unnecessary function calls and stack frames in a hot path, significantly reducing the total runtime.
- **Early termination in Bubble Sort:** Added a `swapped` flag to exit early if no swaps were performed in an inner pass. This prevents unnecessary iterations when the list is already sorted, reducing average case and best case time complexity from always O(n²) to O(n) for already sorted or nearly sorted lists.

**Performance Impact:**
- Reduced hottest-path function call overhead.
- Reduced unnecessary iterations for sorted/permuted lists.
- `swap` in `utils.py` remains for API compatibility but is now unused in the optimized sort.

---""",
        expected_optimized_code="""```python:app/utils.py
def swap(val1, val2):
    # create a helper function for this to count down the number of swaps happened
    # (Unchanged: still a tuple swap for API compatibility, no longer used in sorter)
    temp = val1
    val1 = val2
    val2 = temp
    return (val1, val2)
```
```python:app/main.py
from app.utils import swap


def sorter_test_final10(arr):
    print("codeflash stdout: Sorting list")
    n = len(arr)
    for i in range(n):
        swapped = False  # Optimization: Track if any swaps occurred in this pass
        for j in range(n - 1):
            if arr[j] > arr[j + 1]:
                # Fast tuple swap: avoids function call overhead
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
                swapped = True
        if not swapped:
            break  # Stop if the list is already sorted
    print(f"result: {arr}")
    return arr
```""",
    ),
]


multi_refiner_test_cases: list[RefinerTestCase] = [
    RefinerTestCase(
        optimized_code="""```python:app/utils.py
def swap(val1, val2):
    # Optimized swap using tuple unpacking for better performance
    return val2, val1
```
```python:app/main.py
from app.utils import swap

def sorter_test_final10(arr):
    print("codeflash stdout: Sorting list")
    n = len(arr)
    for i in range(n):
        # Optimized: inner loop range reduces as largest elements "bubble" to end
        for j in range(n - 1 - i):
            if arr[j] > arr[j + 1]:
                # Direct in-place swap, avoids function call overhead
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
    print(f"result: {arr}")
    return arr
```
""",
        llm_response="""Looking at the optimization results, I can see that the main performance improvements come from:

1. **Bubble sort optimization**: Reducing the inner loop range by `i` each iteration (from `range(len(arr) - 1)` to `range(n - 1 - i)`)
2. **Direct tuple swapping**: Eliminating the function call overhead by using direct tuple unpacking

However, I notice some quality issues that can be refined:

1. The `swap` function is no longer being used, but it was modified unnecessarily
2. A new variable `n` was introduced when we could use the original pattern
3. New comments were added that differ from the original style

Let me refine the optimization to be more precise while preserving the core performance improvements:

<replace_in_file>
<path>app/utils.py</path>
<diff>
<<<<<<< SEARCH
def swap(val1, val2):
    # Optimized swap using tuple unpacking for better performance
    return val2, val1
=======
def swap(val1, val2):
    # create a helper function for this to count down the number of swaps happened
    temp = val1
    val1 = val2
    val2 = temp
    return (val1, val2)
>>>>>>> REPLACE

<<<<<<< SEARCH
    # create a helper function for this to count down the number of swaps happened
=======
    # comment is modified
>>>>>>> REPLACE
</diff>
</replace_in_file>

<replace_in_file>
<path>app/main.py</path>
<diff>
<<<<<<< SEARCH
def sorter_test_final10(arr):
    print("codeflash stdout: Sorting list")
    n = len(arr)
    for i in range(n):
        # Optimized: inner loop range reduces as largest elements "bubble" to end
        for j in range(n - 1 - i):
            if arr[j] > arr[j + 1]:
                # Direct in-place swap, avoids function call overhead
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
    print(f"result: {arr}")
    return arr
=======
def sorter_test_final10(arr):
    print("codeflash stdout: Sorting list")
    for i in range(len(arr)):
        for j in range(len(arr) - 1 - i):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
    print(f"result: {arr}")
    return arr
>>>>>>> REPLACE
</diff>
</replace_in_file>

The refinements I made:

1. **Reverted the `swap` function**: Since it's no longer being used in the optimized code, I reverted it to the original implementation to minimize code differences. This doesn't affect performance since the function isn't called.

2. **Removed the `n` variable**: Used the original `len(arr)` pattern to reduce code differences while preserving the key optimization of `range(len(arr) - 1 - i)`.

3. **Removed optimization comments**: Reverted to the original comment style to maintain consistency with the original codebase.

4. **Preserved the core optimizations**:
   - The bubble sort optimization (`range(len(arr) - 1 - i)` instead of `range(len(arr) - 1)`) is maintained
   - The direct tuple swapping (`arr[j], arr[j + 1] = arr[j + 1], arr[j]`) is maintained

These changes make the optimization more precise by reducing unnecessary differences from the original code while preserving all the performance benefits. The optimization should still achieve the same ~115% speedup.""",
        expected_refined_code="""```python:app/utils.py
def swap(val1, val2):
    # comment is modified
    temp = val1
    val1 = val2
    val2 = temp
    return (val1, val2)
```
```python:app/main.py
from app.utils import swap

def sorter_test_final10(arr):
    print("codeflash stdout: Sorting list")
    for i in range(len(arr)):
        for j in range(len(arr) - 1 - i):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
    print(f"result: {arr}")
    return arr
```""",
    )
]


def test_single_optimizer_context() -> None:
    for t in single_optimizer_test_cases:
        ctx = create_optimizer_context("")
        ctx.extract_code_and_explanation_from_llm_res(t.llm_response)
        ctx.parse_and_generate_candidate_schema()

        assert ctx.extracted_code_and_expl is not None
        assert ctx.extracted_code_and_expl.code == t.expected_optimized_code
        assert ctx.extracted_code_and_expl.explanation == t.expected_explanation
        assert ctx.is_valid_code()


def test_multi_optimizer_context() -> None:
    for t in multi_optimizer_test_cases:
        ctx = create_optimizer_context(t.original_code)
        ctx.extract_code_and_explanation_from_llm_res(t.llm_response)
        candidate = ctx.parse_and_generate_candidate_schema()

        assert candidate is not None
        assert candidate.source_code == t.expected_optimized_code
        assert ctx.extracted_code_and_expl is not None
        assert ctx.extracted_code_and_expl.explanation == t.expected_explanation
        assert ctx.is_valid_code()


def test_search_and_replace_optimizer() -> None:
    for t in optimizer_with_search_and_replace_test_cases:
        ctx: BaseOptimizerContext = create_optimizer_context(t.original_code, DiffMethod.SEARCH_AND_REPLACE)
        ctx.extract_code_and_explanation_from_llm_res(t.llm_response)
        candidate = ctx.parse_and_generate_candidate_schema()

        assert candidate is not None
        assert candidate.source_code == t.expected_optimized_code
        assert ctx.extracted_code_and_expl is not None
        assert ctx.extracted_code_and_expl.explanation == t.expected_explanation
        assert ctx.is_valid_code()


def test_single_refiner_context() -> None:
    for t in single_refiner_test_cases:
        ctx = create_refiner_context(t.optimized_code)
        refined_code = ctx.apply_patches_to_optimized_code(t.llm_response)
        assert refined_code == t.expected_refined_code
        assert ctx.is_valid_refinement(refined_code)


def test_multi_refiner_context() -> None:
    for t in multi_refiner_test_cases:
        ctx = create_refiner_context(t.optimized_code)
        refined_code = ctx.apply_patches_to_optimized_code(t.llm_response)

        assert refined_code == t.expected_refined_code
        assert ctx.is_valid_refinement(refined_code)


def test_split_and_group_markdown_code() -> None:
    code = """```python:path/to/file1.py
file1=__name__
```
```python:path/to/file2.py
file2=__name__
```"""

    expected = {"path/to/file1.py": "file1=__name__", "path/to/file2.py": "file2=__name__"}

    assert split_markdown_code(code) == expected
    assert group_code(expected) == code

    code = """```python:relative/path/to/bubble_sort1.py
def sorter1(arr):
    arr.sort()
    return arr
```
```python:relative/path/to/bubble_sort2.py
def sorter2(arr):
    arr.sort()
    return arr
```"""

    expected = {
        "relative/path/to/bubble_sort1.py": """def sorter1(arr):
    arr.sort()
    return arr""",
        "relative/path/to/bubble_sort2.py": """def sorter2(arr):
    arr.sort()
    return arr""",
    }
    assert split_markdown_code(code) == expected
    assert group_code(expected) == code


def test_system_prompt_for_multi_optimizer() -> None:
    ctx = create_optimizer_context(
        "```python:some_file.py\nprint(\"Hello world\")\n```\n```python:some_other_file.py\nprint('hi')\n````\n"
    )

    assert isinstance(ctx, MultiOptimizerContext)
    assert "separate Python module" in ctx.get_system_prompt("")

    # multi-file context with single file is treated as single-file context
    ctx = create_optimizer_context('```python:some_file.py\nprint("Hello world")\n```\n')
    assert isinstance(ctx, SingleOptimizerContext)
    assert "separate Python module" not in ctx.get_system_prompt("")


def test_single_optimizer_context_with_filename_in_code_block() -> None:
    """Test that LLM responses with python:filename.py in code blocks are parsed correctly."""
    llm_response = """I replaced the recursive implementation with an iterative loop.

```python:filename.py
def gcd(a: int, b: int) -> int:
    while b != 0:
        a, b = b, a % b
    return a
```
"""
    ctx = create_optimizer_context("")
    ctx.extract_code_and_explanation_from_llm_res(llm_response)

    expected_code = """def gcd(a: int, b: int) -> int:
    while b != 0:
        a, b = b, a % b
    return a
"""
    assert ctx.extracted_code_and_expl is not None
    assert ctx.extracted_code_and_expl.code == expected_code
    assert ctx.extracted_code_and_expl.explanation == "I replaced the recursive implementation with an iterative loop."
    assert ctx.is_valid_code()


def test_multi_optimizer_partial_file_return() -> None:
    """Test that LLM returning only modified files (not all original files) is handled correctly.

    This tests the fix for trace_id=5e4306f4-a909-48b1-aeb4-932c1b451129 where the optimizer
    received 7 files but LLMs only returned 1-2 files they actually modified. The strict
    validation previously required exact file count matching, causing all candidates to fail.
    """
    # Original code has 3 files
    original_code = """```python:app/utils.py
def helper():
    return "helper"
```
```python:app/main.py
from app.utils import helper

def process(data):
    for i in range(len(data)):
        data[i] = data[i] * 2
    return data
```
```python:app/constants.py
MAX_SIZE = 100
DEFAULT_VALUE = 0
```
"""
    # LLM only returns 1 file (the one it actually modified)
    llm_response = """I optimized the loop using list comprehension for better performance.

```python:app/main.py
from app.utils import helper

def process(data):
    return [x * 2 for x in data]
```
"""
    ctx = create_optimizer_context(original_code)
    assert isinstance(ctx, MultiOptimizerContext)
    ctx.extract_code_and_explanation_from_llm_res(llm_response)

    # Verify extraction worked correctly
    assert ctx.extracted_code_and_expl is not None
    assert (
        ctx.extracted_code_and_expl.code
        == """```python:app/main.py
from app.utils import helper

def process(data):
    return [x * 2 for x in data]
```"""
    )
    assert (
        ctx.extracted_code_and_expl.explanation
        == "I optimized the loop using list comprehension for better performance."
    )

    # Should succeed - partial return is valid
    assert ctx.is_valid_code() is True

    candidate = ctx.parse_and_generate_candidate_schema()
    assert candidate is not None
    assert candidate.optimization_id is not None

    # The returned code should only contain the modified file (with post-processing applied)
    assert (
        candidate.source_code
        == """```python:app/main.py
from app.utils import helper


def process(data):
    return [x * 2 for x in data]
```"""
    )
    assert candidate.explanation == "I optimized the loop using list comprehension for better performance."


def test_multi_optimizer_rejects_new_files_not_in_original() -> None:
    """Test that LLM returning files NOT in the original context is rejected.

    While partial returns (subset of original files) are allowed, returning
    new files that weren't in the original context should be rejected.
    """
    # Need at least 2 files to get a MultiOptimizerContext
    original_code = """```python:app/main.py
def process(data):
    return data * 2
```
```python:app/utils.py
def helper():
    return "helper"
```
"""
    # LLM returns a file that wasn't in the original context
    llm_response = """I added a new helper module.

```python:app/new_helpers.py
def double(x):
    return x * 2
```
"""
    ctx = create_optimizer_context(original_code)
    assert isinstance(ctx, MultiOptimizerContext)
    ctx.extract_code_and_explanation_from_llm_res(llm_response)

    # Verify extraction worked (the code was extracted, just not valid)
    assert ctx.extracted_code_and_expl is not None
    assert (
        ctx.extracted_code_and_expl.code
        == """```python:app/new_helpers.py
def double(x):
    return x * 2
```"""
    )
    assert ctx.extracted_code_and_expl.explanation == "I added a new helper module."

    # Should fail - new file not in original context
    assert ctx.is_valid_code() is False

    # parse_and_generate_candidate_schema should return None for invalid files
    candidate = ctx.parse_and_generate_candidate_schema()
    assert candidate is None
