from core.languages.python.optimizer.context_utils.optimizer_context import BaseOptimizerContext
from core.languages.python.optimizer.diff_patches_utils.diff import DiffMethod
from core.languages.python.optimizer.diff_patches_utils.v4a_diff import (
    apply_update,
    extract_update_sections,
    parse_patch_text,
)
from tests.optimizer.test_context import OptimizerTestCase


def test_a_simple_diff() -> None:
    original_content = """def add(a, b):
    return a + b
def add2(a, b):
    return a + b
"""
    llm_res_with_diff = """
Here is the patch:

*** Begin Patch
*** Update File: file.py
@@ def add(a, b):
-    return a + b
+    return a + b + 2
@@ def add2(a, b):
-    return a + b
+    return a + b + 2
*** End Patch

and some other explanation from the llm.
"""
    file_to_diff = extract_update_sections(llm_res_with_diff)
    update_diff = file_to_diff["file.py"]
    patch = parse_patch_text(update_diff.strip().splitlines(), {"file.py": original_content})
    new_content = apply_update(original_content, patch.actions["file.py"], "file.py")

    assert (
        new_content
        == """def add(a, b):
    return a + b + 2
def add2(a, b):
    return a + b + 2
"""
    )


test_cases: list[OptimizerTestCase] = [
    OptimizerTestCase(
        llm_response="""*** Begin Patch
*** Update File: 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
+    # Optimized bubble sort: stop if no swaps, and shrink inner loop
+    n = len(arr)
+    for i in range(n):
+        swapped = False
+        end = n - i - 1
+        for j in range(end):
+            if arr[j] > arr[j + 1]:
+                arr[j], arr[j + 1] = arr[j + 1], arr[j]
+                swapped = True
+        if not swapped:
+            break
+    return arr
*** End Patch

**Explanation**:

- Used an optimized bubble sort: the inner loop range reduces with each outer loop iteration (`n - i - 1`), since the last `i` elements are sorted.
- Added a `swapped` flag for early termination: if no swaps occur in a pass, the array is already sorted.
- Used tuple assignment to swap elements directly, which is faster and more Pythonic.
- No behavioral changes: function signature, mutation of `arr`, and returned value are all preserved. Comments are kept as required.""",
        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
```
""",
        expected_explanation="""**Explanation**:

- Used an optimized bubble sort: the inner loop range reduces with each outer loop iteration (`n - i - 1`), since the last `i` elements are sorted.
- Added a `swapped` flag for early termination: if no swaps occur in a pass, the array is already sorted.
- Used tuple assignment to swap elements directly, which is faster and more Pythonic.
- No behavioral changes: function signature, mutation of `arr`, and returned value are all preserved. Comments are kept as required.""",
        expected_optimized_code="""```python:app/sorter.py
def sorter(arr):
    # Optimized bubble sort: stop if no swaps, and shrink inner loop
    n = len(arr)
    for i in range(n):
        swapped = False
        end = n - i - 1
        for j in range(end):
            if arr[j] > arr[j + 1]:
                arr[j], arr[j + 1] = arr[j + 1], arr[j]
                swapped = True
        if not swapped:
            break
    return arr
```""",
    ),
    OptimizerTestCase(
        llm_response="""*** Begin Patch
*** Update File: app/file1.py
@@
*** Update File: app/file2.py
@@
*** End Patch

No changes.""",
        original_code="""```python:app/file1.py
def say_hi():
    print('hi')
```

```python:app/file2.py
def say_hi_2():
    print('hi')
```
""",
        expected_explanation="""No changes.""",
        expected_optimized_code="""```python:app/file1.py
def say_hi():
    print('hi')
```

```python:app/file2.py
def say_hi_2():
    print('hi')
```
""",
    ),
    OptimizerTestCase(
        llm_response="""*** Begin Patch
*** Update File: app/file1.py
@@ def say_hi():
-    print('hi')
+    print('hello')
*** Update File: app/file2.py
@@ def say_hi_2():
-    print('hi')
+    print('hello')
*** End Patch

**Explanation**:
this is some explanation.""",
        original_code="""```python:app/file1.py
def say_hi():
    print('hi')
```

```python:app/file2.py
def say_hi_2():
    print('hi')
```
""",
        expected_explanation="""**Explanation**:
this is some explanation.""",
        expected_optimized_code="""```python:app/file1.py
def say_hi():
    print('hello')
```
```python:app/file2.py
def say_hi_2():
    print('hello')
```""",
    ),
]


def test_v4a_diff() -> None:
    for t in test_cases:
        ctx = BaseOptimizerContext.get_dynamic_context("", "", t.original_code, DiffMethod.V4A)

        ctx.extract_code_and_explanation_from_llm_res(t.llm_response)
        candidate = ctx.parse_and_generate_candidate_schema()

        if not candidate:
            # candidate is not accepted
            assert t.expected_optimized_code == t.original_code
            assert ctx.extracted_code_and_expl is not None
            assert t.expected_explanation == ctx.extracted_code_and_expl.explanation
            continue

        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()
