from aiservice.models.functions_to_optimize import FunctionParent, FunctionToOptimize

from testgen.instrumentation.edit_generated_test import parse_module_to_cst
from testgen.postprocessing.removeassert_transformer import remove_asserts_from_test


def test_remove_asserts() -> None:
    original_test = """# imports
import pytest  # used for our unit tests
from code_to_optimize.bubble_sort import sorter

# unit tests

def test_basic_functionality():
    # Simple Unsorted List
    res = sorter([3, 1, 2])
    assert res == [1, 2, 3]
    res = sorter([5, 3, 8, 4, 2])
    assert res == [2, 3, 4, 5, 8]
    # Already Sorted List
    assert sorter([1, 2, 3]) == [1, 2, 3]
    assert sorter([2, 4, 6, 8]) == [2, 4, 6, 8]
    # Reverse Sorted List
    assert sorter([3, 2, 1]) == [1, 2, 3]
    assert sorter([9, 7, 5, 3, 1]) == [1, 3, 5, 7, 9]
    assert [1,2,3,4,5] == sorter([1,2,3,4,5])
    with open("file.txt", "w") as f:
        assert sorter([1, 2, 3, 4, 5]) == [1, 2, 3, 4, 5]
"""

    expected = """# imports
from code_to_optimize.bubble_sort import sorter
import pytest  # used for our unit tests

# unit tests

def test_basic_functionality():
    # Simple Unsorted List
    codeflash_output = sorter([3, 1, 2])
    codeflash_output = sorter([5, 3, 8, 4, 2])
    # Already Sorted List
    codeflash_output = sorter([1, 2, 3])
    codeflash_output = sorter([2, 4, 6, 8])
    # Reverse Sorted List
    codeflash_output = sorter([3, 2, 1])
    codeflash_output = sorter([9, 7, 5, 3, 1])
    codeflash_output = sorter([1,2,3,4,5])
    with open("file.txt", "w") as f:
        codeflash_output = sorter([1, 2, 3, 4, 5])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter", file_path="bubble_sort.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert result == expected


def test_remove_asserts_class_functions() -> None:
    original_test = """# imports
import pytest  # used for our unit tests
from code_to_optimize.bubble_sort import BubbleSortClass

bubble_sort = BubbleSortClass()

# unit tests

def test_basic_functionality():
    # Simple Unsorted List
    res = bubble_sort.sorter([3, 1, 2])
    assert res == [1, 2, 3]
    res = bubble_sort.sorter([5, 3, 8, 4, 2])
    assert res == [2, 3, 4, 5, 8]
    # Already Sorted List
    assert bubble_sort.sorter([1, 2, 3]) == [1, 2, 3]
    assert bubble_sort.sorter([2, 4, 6, 8]) == [2, 4, 6, 8]"""

    expected = """# imports
from code_to_optimize.bubble_sort import BubbleSortClass
import pytest  # used for our unit tests

bubble_sort = BubbleSortClass()

# unit tests

def test_basic_functionality():
    # Simple Unsorted List
    codeflash_output = bubble_sort.sorter([3, 1, 2])
    codeflash_output = bubble_sort.sorter([5, 3, 8, 4, 2])
    # Already Sorted List
    codeflash_output = bubble_sort.sorter([1, 2, 3])
    codeflash_output = bubble_sort.sorter([2, 4, 6, 8])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter",
        file_path="bubble_sort.py",
        parents=[FunctionParent(name="BubbleSortClass", type="ClassDef")],
        starting_line=None,
        ending_line=None,
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert result == expected


def test_remove_asserts_unittest() -> None:
    original_test = """import unittest

from code_to_optimize.bubble_sort import sorter


class TestPigLatin(unittest.TestCase):
    def test_sort(self):
        input = [5, 4, 3, 2, 1, 0]
        self.assertEqual(sorter(input), [0, 1, 2, 3, 4, 5])

        input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
        output = sorter(input)
        self.assertEqual(output, [0.0, 1.0, 2.0, 3.0, 4.0, 5.0])

        input = list(reversed(range(5000)))
        output = sorter(input)
        self.assertEqual(output, list(range(5000)))
        self.assertEqual(output, sorter(input))
        with open("file.txt", "w") as f:
            self.assertEqual(sorter(input), list(range(5000)))
"""

    expected = """from code_to_optimize.bubble_sort import sorter
import unittest


class TestPigLatin(unittest.TestCase):
    def test_sort(self):
        input = [5, 4, 3, 2, 1, 0]
        codeflash_output = sorter(input)

        input = [5.0, 4.0, 3.0, 2.0, 1.0, 0.0]
        codeflash_output = sorter(input)

        input = list(reversed(range(5000)))
        codeflash_output = sorter(input)
        codeflash_output = sorter(input)
        with open("file.txt", "w") as f:
            codeflash_output = sorter(input)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter", file_path="bubble_sort.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert result == expected


def test_remove_single_arg_asserts_unittest() -> None:
    original_test = """import unittest
from code_to_optimize.bubble_sort import sorter

class TestBubbleSort(unittest.TestCase):
    def test_sort(self):
        input = [5, 4, 3, 2, 1]
        output = sorter(input)
        self.assertTrue(output)
        self.assertFalse(sorter([]))

        input = [3, 2, 1]
        self.assertTrue(sorter(input))
        self.assertIsNone(sorter([]))
        self.assertIsNotNone(sorter([1]))
"""

    expected = """from code_to_optimize.bubble_sort import sorter
import unittest

class TestBubbleSort(unittest.TestCase):
    def test_sort(self):
        input = [5, 4, 3, 2, 1]
        codeflash_output = sorter(input)
        codeflash_output = sorter([])

        input = [3, 2, 1]
        codeflash_output = sorter(input)
        codeflash_output = sorter([])
        codeflash_output = sorter([1])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter", file_path="bubble_sort.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert result == expected


def test_remove_irrelevant_assert() -> None:
    original_test = """def test_relative_validity_no_tree():
    hdbscan = HDBSCAN()
    assert hdbscan.relative_validity_() == 0.5
    result = 5
    assert result == 5
    assert 1 == True"""

    expected = """from code_to_optimize.bubble_sort import HDBSCAN
def test_relative_validity_no_tree():
    hdbscan = HDBSCAN()
    codeflash_output = hdbscan.relative_validity_()
    result = 5
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""

    function_to_optimize = FunctionToOptimize(
        function_name="relative_validity_",
        file_path="/tmp/path",
        parents=[FunctionParent(name="HDBSCAN", type="ClassDef")],
    )

    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert result == expected


def test_remove_asserts_one_line_functions() -> None:
    original_test = """def test_short_function(): x = 5; y = sorter([3,2,1]); assert y == [1,2,3]; z = 10"""

    expected = """from code_to_optimize.bubble_sort import sorter
def test_short_function(): x = 5; codeflash_output = sorter([3,2,1]); z = 10
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter", file_path="bubble_sort.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert result == expected


def test_remove_asserts_one_line_statements() -> None:
    original_test = """def test_short_function():
    x = 5; y = sorter([3,2,1]); assert y == [1,2,3]; z = 10
    a = sorter([5,4,3]); assert a == [3,4,5]"""

    expected = """from code_to_optimize.bubble_sort import sorter
def test_short_function():
    x = 5; codeflash_output = sorter([3,2,1]); z = 10
    codeflash_output = sorter([5,4,3])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter", file_path="bubble_sort.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert result == expected


def test_remove_asserts_mixed_statement_types() -> None:
    original_test = """def test_mixed_statements():
    # First part has one-line statements
    x = 5; y = sorter([3,2,1]); assert y == [1,2,3]

    # Second part has regular statements
    codeflash_output = sorter([5,4,3])
    assert result == [3,4,5]"""

    expected = """from code_to_optimize.bubble_sort import sorter
def test_mixed_statements():
    # First part has one-line statements
    x = 5; codeflash_output = sorter([3,2,1])

    # Second part has regular statements
    codeflash_output = sorter([5,4,3])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter", file_path="bubble_sort.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert result == expected


def test_function_with_pass() -> None:
    original_test = """def test_large_callable_object():
    class LargeClass:
        async def __call__(self):
            pass
        def method1(self): pass
        def method2(self): pass
        # Add more methods if needed to simulate a large class
    assert is_coroutine_callable(LargeClass()) == True"""
    expected = """from some_file import is_coroutine_callable
def test_large_callable_object():
    class LargeClass:
        async def __call__(self):
            pass
        def method1(self): pass
        def method2(self): pass
        # Add more methods if needed to simulate a large class
    codeflash_output = is_coroutine_callable(LargeClass())
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""
    function_to_optimize = FunctionToOptimize(
        function_name="is_coroutine_callable",
        file_path="some_file.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="some_file",
    ).code
    assert result == expected


def test_assert_not_with_primitive() -> None:
    original_test = """def test_primitive_types():
    assert not is_coroutine_callable(42)
    assert not is_coroutine_callable("string")
    assert not is_coroutine_callable([1, 2, 3])"""
    expected = """from some_file import is_coroutine_callable
def test_primitive_types():
    codeflash_output = not is_coroutine_callable(42)
    codeflash_output = not is_coroutine_callable("string")
    codeflash_output = not is_coroutine_callable([1, 2, 3])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""
    function_to_optimize = FunctionToOptimize(
        function_name="is_coroutine_callable",
        file_path="some_file.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="some_file",
    ).code
    assert result == expected


def test_assert_with_primitive() -> None:
    original_test = """def test_primitive_types():
    assert is_coroutine_callable(42)
    assert is_coroutine_callable("string")
    assert is_coroutine_callable([1, 2, 3])"""
    expected = """from some_file import is_coroutine_callable
def test_primitive_types():
    codeflash_output = is_coroutine_callable(42)
    codeflash_output = is_coroutine_callable("string")
    codeflash_output = is_coroutine_callable([1, 2, 3])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""
    function_to_optimize = FunctionToOptimize(
        function_name="is_coroutine_callable",
        file_path="some_file.py",
        parents=[],
        starting_line=None,
        ending_line=None,
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="some_file",
    ).code
    assert result == expected


def test_function_for_list_comprehension() -> None:
    original_test = """def test_list_comprehension():
    assert [some_function(x) for x in range(5)] == [0, 1, 2, 3, 4]"""
    expected = """from some_file import some_function
def test_list_comprehension():
    codeflash_output = [some_function(x) for x in range(5)]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""
    function_to_optimize = FunctionToOptimize(
        function_name="some_function", file_path="some_file.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="some_file",
    ).code
    assert result == expected


def test_singleline_function_for_list_comprehension() -> None:
    original_test = """def test_list_comprehension(): assert [some_function(x) for x in range(5)] == [0, 1, 2, 3, 4]; assert [some_function(y) for x in range(4)] == [1,3,4,5]"""
    expected = """from some_file import some_function
def test_list_comprehension(): codeflash_output = [some_function(x) for x in range(5)]; codeflash_output = [some_function(y) for x in range(4)]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""
    function_to_optimize = FunctionToOptimize(
        function_name="some_function", file_path="some_file.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="some_file",
    ).code
    assert result == expected


def test_unittest_function_for_list_comprehension() -> None:
    original_test = """def test_unittest_style():
    self.assertEqual([some_function(x) for x in range(5)], [0, 1, 2, 3, 4])"""
    expected = """from some_file import some_function
def test_unittest_style():
    codeflash_output = [some_function(x) for x in range(5)]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code."""
    function_to_optimize = FunctionToOptimize(
        function_name="some_function", file_path="some_file.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="some_file",
    ).code
    assert result == expected


def test_remove_asserts_import_star() -> None:
    original_test = """from x import *
from code_to_optimize.bubble_sort import sorter
from code_to_optimize.bubble_sort import *
from y import *

def test_basic_functionality():
    assert sorter([3, 1, 2]) == [1, 2, 3]
"""

    expected = """from code_to_optimize.bubble_sort import sorter
from x import *
from code_to_optimize.bubble_sort import *
from y import *

def test_basic_functionality():
    codeflash_output = sorter([3, 1, 2])
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
"""

    function_to_optimize = FunctionToOptimize(
        function_name="sorter", file_path="bubble_sort.py", parents=[], starting_line=None, ending_line=None
    )
    result = remove_asserts_from_test(
        module=parse_module_to_cst(original_test),
        function_to_optimize=function_to_optimize,
        helper_function_names=[],
        module_path="code_to_optimize.bubble_sort",
    ).code
    assert result == expected
