import libcst
import libcst as cst

from core.languages.python.optimizer.models import CodeExplanationAndID
from core.languages.python.optimizer.postprocess import (
    cleanup_explanations,
    dedup_and_sort_imports,
    deduplicate_optimizations,
    equality_check,
    filter_ellipsis_containing_code,
    fix_missing_docstring,
    optimizations_postprocessing_pipeline,
)


def test_cleanup_explanations_removes_code_markers() -> None:
    original_module = cst.parse_module("def example(): pass")
    empty_cst_module = cst.parse_module("")
    explanations = [
        CodeExplanationAndID(empty_cst_module, "Hi! Here is the code:\n```python\ndef example(): pass\n```", "1"),
        CodeExplanationAndID(empty_cst_module, "```python\ndef example(): pass\n```\nEnd of explanation.", "2"),
    ]
    expected_explanations = [
        CodeExplanationAndID(empty_cst_module, "Hi!\n", "1"),
        CodeExplanationAndID(empty_cst_module, "\nEnd of explanation.", "2"),
    ]
    cleaned = cleanup_explanations(original_module, explanations)
    assert cleaned == expected_explanations


def test_cleanup_explanations_preserves_non_code_text() -> None:
    original_module = cst.parse_module("def example(): pass")
    empty_cst_module = cst.parse_module("")
    explanations = [
        CodeExplanationAndID(
            empty_cst_module, "This is a test. Here is the code:\n```python\ndef example(): pass\n```", "1"
        ),
        CodeExplanationAndID(
            empty_cst_module,
            "Explanation before code.\n```python\ndef example(): pass\n```\nExplanation after code.",
            "2",
        ),
    ]
    expected_explanations = [
        CodeExplanationAndID(empty_cst_module, "This is a test.\n", "1"),
        CodeExplanationAndID(empty_cst_module, "Explanation before code.\n\nExplanation after code.", "2"),
    ]
    cleaned = cleanup_explanations(original_module, explanations)
    assert cleaned == expected_explanations


def test_postprocess_optimizations_basic() -> None:
    original_module = cst.parse_module("print('hello')")
    optimizations = [CodeExplanationAndID(libcst.parse_module("print('hi')"), "Simplified print statement", "1")]

    expected = [CodeExplanationAndID(libcst.parse_module("print('hi')"), "Simplified print statement", "1")]
    actual = optimizations_postprocessing_pipeline(original_module, optimizations)

    assert actual[0].cst_module.deep_equals(expected[0].cst_module)


def test_postprocess_deduplicates() -> None:
    original_module = cst.parse_module("print('hello')")
    optimizations = [
        CodeExplanationAndID(libcst.parse_module("print('hi')"), "Simplified print", "1"),
        CodeExplanationAndID(libcst.parse_module("print('hi')"), "Simplified print again", "2"),
    ]

    expected = [CodeExplanationAndID(libcst.parse_module("print('hi')"), "Simplified print again", "2")]
    actual = deduplicate_optimizations(original_module, optimizations)

    assert actual[0].cst_module.deep_equals(expected[0].cst_module)


def test_equality_check() -> None:
    original_code = '''from nuitka.utils.SlotMetaClasses import getMetaClassBase

class NuitkaPluginBase(getMetaClassBase("Plugin", require_slots=False)):
    def isRequiredImplicitImport(self, module, full_name):
        """Indicate whether an implicitly imported module should be accepted.

        Notes:
            You may negate importing a module specified as "implicit import",
            although this is an unexpected event.

        Args:
            module: the module object
            full_name: of the implicitly import module
        Returns:
            True or False
        """
        # Virtual method, pylint: disable=no-self-use,unused-argument
        return True
    '''
    original_module = cst.parse_module(original_code)
    optimizations = [
        CodeExplanationAndID(
            libcst.parse_module(
                '''from nuitka.utils.SlotMetaClasses import getMetaClassBase

class NuitkaPluginBase(getMetaClassBase("Plugin", require_slots=False)):
    def isRequiredImplicitImport(self, module, full_name):
        """Indicate whether an implicitly imported module should be accepted.

        Notes:
            You may negate importing a module specified as "implicit import",
            although this is an unexpected event.

        Args:
            module: the module object
            full_name: of the implicitly import module
        Returns:
            True or False
        """
        # Return True for all implicitly imported modules
        return True
'''
            ),
            "Simplified print",
            "1",
        )
    ]
    actual = equality_check(original_module, optimizations)
    assert len(actual) == 0


def test_postprocess_bubble_sort() -> None:
    original_code = """
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
"""
    original_module = cst.parse_module(original_code)

    optimizations = [
        CodeExplanationAndID(
            libcst.parse_module("def sorter(arr):\n    return sorted(arr)\n"),
            "Optimized to use built-in sorted() function",
            "1",
        )
    ]

    expected = [
        CodeExplanationAndID(
            libcst.parse_module("\ndef sorter(arr):\n    return sorted(arr)\n"),
            "Optimized to use built-in sorted() function",
            "1",
        )
    ]

    actual = optimizations_postprocessing_pipeline(original_module, optimizations)

    assert actual[0].cst_module.deep_equals(expected[0].cst_module)


def test_optimize_bubble_sort_multiple() -> None:
    original_code = """
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
"""
    original_module = cst.parse_module(original_code)

    optimizations = [
        CodeExplanationAndID(
            libcst.parse_module("def sorter(arr):\n return sorted(arr)\n"),
            "Your program is performing a bubble sort, which is known for its poor performance in large datasets because it has a worst and average time complexity of O(n^2). We can make this run faster by using a more efficient sorting algorithm. Python's built-in sorted() function uses the Timsort algorithm, which has an average and worst case time complexity of O(n log n), making it much faster and more efficient for large datasets.\n\nHere's how you can rewrite your function:\n\n\nThis change significantly improves the time complexity and memory usage of your function.",
            "1",
        ),
        CodeExplanationAndID(
            libcst.parse_module("def sorter(arr):\n return sorted(arr)\n"),
            "Your original program is using bubble sort which has a time complexity of O(n^2) making it inefficient for large data sets. A faster sorting algorithm could be applied here, 'Timsort', which is a hybrid sorting algorithm, derived from merge sort and insertion sort, designed to perform well on many kinds of real-world data. This algorithm is built-in as a native sorting algorithm in Python, and hence it runs considerably faster than any algorithm implemented in Python.\n\nHere is the updated faster program:\n\n",
            "2",
        ),
        CodeExplanationAndID(libcst.parse_module(original_code), "Blah blah Blah.", "3"),
    ]

    expected = [
        CodeExplanationAndID(
            libcst.parse_module("\ndef sorter(arr):\n return sorted(arr)\n"),
            "Your original program is using bubble sort which has a time complexity of O(n^2) making it inefficient for large data sets. A faster sorting algorithm could be applied here, 'Timsort', which is a hybrid sorting algorithm, derived from merge sort and insertion sort, designed to perform well on many kinds of real-world data. This algorithm is built-in as a native sorting algorithm in Python, and hence it runs considerably faster than any algorithm implemented in Python.\n\nHere is the updated faster program.\n\n",
            "2",
        )
    ]

    actual = optimizations_postprocessing_pipeline(original_module, optimizations)

    assert actual[0].cst_module.deep_equals(expected[0].cst_module)


def test_filter_explanations() -> None:
    empty_cst_module = cst.parse_module("")

    explanations = [
        """Here is an optimized version of the function. I used an compiled regular expression for better performance and file.readlines() for faster file reading. Here is the code:

This code could be even faster for larger files as it reads the file all at once and uses the compiled regular expression to match lines. It also reduces function calls in the loop, which tend to be expensive, and uses native Python operations.""",
        """Firstly, I noticed some repeated code which you can eliminate. This can improve the runtime performance of your application and also the readability of your code. In the two functions is_target_function_line() and update_line_node(), you loop through ast.walk(line_node) twice to look for a matched function name. This operation is unnecessary and can be improved to iterate once.

I provide an optimized version of your code below. I've used a dictionary to store the result of ast.walk(line_node) so that we only need to scan the nodes once. Please note that refactoring might raise unexpected side effects, depending on the interlink of these functions with other parts of the code. It's advisable to thoroughly test your application after changes.""",
        """The provided program implements a bubble sort, which has a time complexity of O(n^2). A more efficient sorting algorithm is Timsort, which Python's built-in sorted function uses. It has a time complexity of O(n log n). Here's how you can rewrite your function using Python's built-in sorting function:
```
def sorter(arr):
    return sorted(arr)
```
The built-in sorted function returns a new sorted list, and does not modify the original list. If you want to sort the list in place (modify the original list), you can use the sort method of list objects:

Other than the notable performance difference with large input data due to the more effective sorting algorithm, these function versions also incorporate Python's swapping mechanism which too saves additional time compared to traditional element swapping.""",
        """Your given python function is already optimized to the maximum, there's no possible way to rewrite it in a faster way. It simply returns a value from a dictionary which is an O(1) operation meaning it will take the same constant time no matter how large the input dictionary is.

This function is performing in constant time (O(1)) and can't be optimized further.""",
        """The original code you provided is already quite efficient in terms of computational complexity. Nonetheless, we can still make a small optimization to it. Instead of using str.join to turn a list of commands into a string and executing it with subprocess.run, we can pass the list directly to subprocess.run, this way, python will handle it optimally.

Here is the optimized code:

The changes made:

    Running the command directly from the list hf_spaces_deploy_cmd.
    I added str(e) in the except block. This is because formatting a low-level exception may cause another exception because not all exceptions implicitly call str() when they're in a format string. So it's safer to do it explicitly.

These changes are mostly about making the code a little bit more Pythonic and safer, they may not notably speed up the execution. Your original code is already written well and does not do any heavy computations to optimize.""",
        """Here is an optimized version of the function. I used an compiled regular expression for better performance and file.readlines() for faster file reading. Here is the code:

This code could be even faster for larger files as it reads the file all at once and uses the compiled regular expression to match lines. It also reduces function calls in the loop, which tend to be expensive, and uses native Python operations.""",
        """The current program is using a try-except block and calling the logging.warning function when the length of the string is zero. However, this does not need to be caught as an exception and can be handled in a simple if condition check. Also, the in operator is inefficient as it scans through all the characters in string.printable for every character in s resulting in a time complexity of O(n*m). A one-time creation of set(string.printable) can greatly improve the execution time. Here is the optimized version of the code:

This version of the program should run faster for larger inputs. Note that this program makes a trade-off between runtime and memory, using additional memory for storing string.printable in a set.""",
        """The provided code is already well optimized in terms of Python, as it uses the dictionary's get() method which provides a default value when a specified key is not in the dictionary and avoids raising an exception. This is more efficient than checking if the key is in the dictionary then getting the value. However, a slight enhancement would be pre-computing default values in the else statement:

This may slightly speed up runtime due to reduction in function calls but any performance gain would most likely be negligible as the given code was already quite efficient. Also, remember that optimizations can sometimes make code harder to read and should only be pursued when necessary.""",
        """However, if the default parameter config = None is a common case and creating a new instance of BaseEmbedderConfig is a heavy or costly operation, we could improve it by using a global mutable default parameter to avoid creating a new instance each time we call __init__, as follows:

Note: This suggested optimization should be applied carefully because if you modify the default_config object, modifications will persist across calls which can lead to unexpected behavior. If BaseEmbedderConfig's instances are not supposed to be modified after creation, this would not be an issue.

The rest of your code looks efficient and there is not much opportunity for further optimization given the current information. However, keep in mind that the performance of the __init__ method and to_embeddings heavily depends on the BaseEmbedderConfig class and embedding_fn method, which are not provided. Those might provide opportunities for optimization.""",
        """The given code mostly consists of class and method definitions, so there's limited room for optimization because the computational complexity is relatively low. However, there are minor changes that can be applied for a bit of efficiency:

Using is not None comparisons in place of or checks in the constructor could prevent potential issues with zero values (0, '', False, etc.), which are False in a boolean context.

Used c.__class__ is MessageContentText instead of isinstance(c, MessageContentText). This checks for the exact type, not subclasses, which is faster if no subclassing is being used.""",
        """The Python program you've provided isn't necessarily inefficient. However, certain parts of the code contain redundant or repetitive operations, or they could benefit from pre-computing some values. Here are improved versions of these functions.

Most changes in the code are associated with avoiding repetitive calls to certain functions, storing the results, and reusing them if needed. This can make the program run faster.""",
    ]

    expected = [
        """Here is an optimized version of the function. I used an compiled regular expression for better performance and file.readlines() for faster file reading.

This code could be even faster for larger files as it reads the file all at once and uses the compiled regular expression to match lines. It also reduces function calls in the loop, which tend to be expensive, and uses native Python operations.""",
        """Firstly, I noticed some repeated code which you can eliminate. This can improve the runtime performance of your application and also the readability of your code. In the two functions is_target_function_line() and update_line_node(), you loop through ast.walk(line_node) twice to look for a matched function name. This operation is unnecessary and can be improved to iterate once.

I provide an optimized version of your code below. I've used a dictionary to store the result of ast.walk(line_node) so that we only need to scan the nodes once. Please note that refactoring might raise unexpected side effects, depending on the interlink of these functions with other parts of the code. It's advisable to thoroughly test your application after changes.""",
        """The provided program implements a bubble sort, which has a time complexity of O(n^2). A more efficient sorting algorithm is Timsort, which Python's built-in sorted function uses. It has a time complexity of O(n log n). Here's how you can rewrite your function using Python's built-in sorting function.

The built-in sorted function returns a new sorted list, and does not modify the original list. If you want to sort the list in place (modify the original list), you can use the sort method of list objects.

Other than the notable performance difference with large input data due to the more effective sorting algorithm, these function versions also incorporate Python's swapping mechanism which too saves additional time compared to traditional element swapping.""",
        """Your given python function is already optimized to the maximum, there's no possible way to rewrite it in a faster way. It simply returns a value from a dictionary which is an O(1) operation meaning it will take the same constant time no matter how large the input dictionary is.

This function is performing in constant time (O(1)) and can't be optimized further.""",
        """The original code you provided is already quite efficient in terms of computational complexity. Nonetheless, we can still make a small optimization to it. Instead of using str.join to turn a list of commands into a string and executing it with subprocess.run, we can pass the list directly to subprocess.run, this way, python will handle it optimally.


The changes made.

    Running the command directly from the list hf_spaces_deploy_cmd.
    I added str(e) in the except block. This is because formatting a low-level exception may cause another exception because not all exceptions implicitly call str() when they're in a format string. So it's safer to do it explicitly.

These changes are mostly about making the code a little bit more Pythonic and safer, they may not notably speed up the execution. Your original code is already written well and does not do any heavy computations to optimize.""",
        """Here is an optimized version of the function. I used an compiled regular expression for better performance and file.readlines() for faster file reading.

This code could be even faster for larger files as it reads the file all at once and uses the compiled regular expression to match lines. It also reduces function calls in the loop, which tend to be expensive, and uses native Python operations.""",
        """The current program is using a try-except block and calling the logging.warning function when the length of the string is zero. However, this does not need to be caught as an exception and can be handled in a simple if condition check. Also, the in operator is inefficient as it scans through all the characters in string.printable for every character in s resulting in a time complexity of O(n*m). A one-time creation of set(string.printable) can greatly improve the execution time.

This version of the program should run faster for larger inputs. Note that this program makes a trade-off between runtime and memory, using additional memory for storing string.printable in a set.""",
        """The provided code is already well optimized in terms of Python, as it uses the dictionary's get() method which provides a default value when a specified key is not in the dictionary and avoids raising an exception. This is more efficient than checking if the key is in the dictionary then getting the value. However, a slight enhancement would be pre-computing default values in the else statement.

This may slightly speed up runtime due to reduction in function calls but any performance gain would most likely be negligible as the given code was already quite efficient. Also, remember that optimizations can sometimes make code harder to read and should only be pursued when necessary.""",
        """However, if the default parameter config = None is a common case and creating a new instance of BaseEmbedderConfig is a heavy or costly operation, we could improve it by using a global mutable default parameter to avoid creating a new instance each time we call __init__.

Note: This suggested optimization should be applied carefully because if you modify the default_config object, modifications will persist across calls which can lead to unexpected behavior. If BaseEmbedderConfig's instances are not supposed to be modified after creation, this would not be an issue.

The rest of your code looks efficient and there is not much opportunity for further optimization given the current information. However, keep in mind that the performance of the __init__ method and to_embeddings heavily depends on the BaseEmbedderConfig class and embedding_fn method, which are not provided. Those might provide opportunities for optimization.""",
        """The given code mostly consists of class and method definitions, so there's limited room for optimization because the computational complexity is relatively low. However, there are minor changes that can be applied for a bit of efficiency.

Using is not None comparisons in place of or checks in the constructor could prevent potential issues with zero values (0, '', False, etc.), which are False in a boolean context.

Used c.__class__ is MessageContentText instead of isinstance(c, MessageContentText). This checks for the exact type, not subclasses, which is faster if no subclassing is being used.""",
        """The Python program you've provided isn't necessarily inefficient. However, certain parts of the code contain redundant or repetitive operations, or they could benefit from pre-computing some values.

Most changes in the code are associated with avoiding repetitive calls to certain functions, storing the results, and reusing them if needed. This can make the program run faster.""",
    ]
    code_and_explanations = [CodeExplanationAndID(empty_cst_module, exp, str(i)) for i, exp in enumerate(explanations)]
    expected_code_and_explanations = [
        CodeExplanationAndID(empty_cst_module, exp, str(i)) for i, exp in enumerate(expected)
    ]
    filtered = cleanup_explanations(cst.parse_module(""), code_and_explanations)
    for filtered_elem, expected_elem in zip(filtered, expected_code_and_explanations, strict=False):
        assert filtered_elem.explanation == expected_elem.explanation


def test_docstring_fix() -> None:
    original_code = '''def test_1():
    """useful docstring"""
    pass
'''
    original_module = cst.parse_module(original_code)

    optimized_code = libcst.parse_module(
        """def test_1():
    pass
"""
    )
    expected = libcst.parse_module(
        '''def test_1():
    """useful docstring"""
    pass
'''
    )
    actual = fix_missing_docstring(original_module, [CodeExplanationAndID(optimized_code, "Explanation", "1")])

    assert actual[0].cst_module.deep_equals(expected)

    original_code = '''def test_1():
    """useful docstring
    has a lot of multi-line details."""
    pass
def test_2():
    """useful docstring"""
    pass
    '''
    original_module = cst.parse_module(original_code)
    optimized_code = libcst.parse_module(
        '''def test_1():
    pass
def test_2():
    """useful docstring v2"""
    pass
    '''
    )
    expected = libcst.parse_module(
        '''def test_1():
    """useful docstring
    has a lot of multi-line details."""
    pass
def test_2():
    """useful docstring"""
    pass'''
    )
    actual = fix_missing_docstring(original_module, [CodeExplanationAndID(optimized_code, "Explanation", "2")])
    assert actual[0].cst_module.deep_equals(expected)

    original_code = '''class TestClass:
    def test_1(self):
        """useful docstring"""
        pass
'''
    original_module = cst.parse_module(original_code)
    optimized_code = libcst.parse_module(
        """class TestClass:
    def test_1(self):
        pass
"""
    )
    expected = libcst.parse_module(
        '''class TestClass:
    def test_1(self):
        """useful docstring"""
        pass
'''
    )
    actual = fix_missing_docstring(original_module, [CodeExplanationAndID(optimized_code, "Explanation", "3")])
    assert actual[0].cst_module.deep_equals(expected)

    original_code = '''def test_1():
    """useful docstring2"""
    pass
class TestClass:
    def test_1(self):
        """useful docstring"""
        pass
    '''
    original_module = cst.parse_module(original_code)
    optimized_code = libcst.parse_module(
        """def test_1():
    pass
class TestClass:
    def test_1(self):
        pass
    """
    )
    expected = libcst.parse_module(
        '''def test_1():
    """useful docstring2"""
    pass
class TestClass:
    def test_1(self):
        """useful docstring"""
        pass
    '''
    )
    actual = fix_missing_docstring(original_module, [CodeExplanationAndID(optimized_code, "Explanation", "4")])
    assert actual[0].cst_module.deep_equals(expected)

    original_code = '''@deprecated("0.1.0", alternative="create_react_agent", removal="0.2.0")
class ZeroShotAgent(Agent):
    """Agent for the MRKL chain."""
    @classmethod
    def create_prompt(
        cls,
        tools: Sequence[BaseTool],
        prefix: str = PREFIX,
        suffix: str = SUFFIX,
        format_instructions: str = FORMAT_INSTRUCTIONS,
        input_variables: Optional[List[str]] = None,
    ) -> PromptTemplate:
        """Create prompt in the style of the zero shot agent.
        Args:
            tools: List of tools the agent will have access to, used to format the
                prompt.
            prefix: String to put before the list of tools.
            suffix: String to put after the list of tools.
            input_variables: List of input variables the final prompt will expect.

        Returns:
            A PromptTemplate with the template assembled from the pieces here.
        """
        pass
'''
    original_module = cst.parse_module(original_code)
    optimized_code = libcst.parse_module(
        """@deprecated("0.1.0", alternative="create_react_agent", removal="0.2.0")
class ZeroShotAgent(Agent):
    @classmethod
    def create_prompt(
        cls,
        tools: Sequence[BaseTool],
        prefix: str = PREFIX,
        suffix: str = SUFFIX,
        format_instructions: str = FORMAT_INSTRUCTIONS,
        input_variables: Optional[List[str]] = None,
    ) -> PromptTemplate:
        pass
"""
    )
    actual = fix_missing_docstring(original_module, [CodeExplanationAndID(optimized_code, "Explanation", "5")])
    assert actual[0].cst_module.deep_equals(libcst.parse_module(original_code))

    original_code = '''def test_1():
    """original has it\n multi line"""
    pass
def test_2():
    pass
def test_3():
    """both have\n original"""
    pass'''
    original_module = cst.parse_module(original_code)
    optimized_code = libcst.parse_module(
        '''def test_1():
    pass
def test_2():
    """retain this one v2\n multi line"""
    pass
def test_3():
    """both have\n original v2"""
    pass'''
    )
    expected = libcst.parse_module(
        '''def test_1():
    """original has it\n multi line"""
    pass
def test_2():
    """retain this one v2\n multi line"""
    pass
def test_3():
    """both have\n original"""
    pass'''
    )
    actual = fix_missing_docstring(original_module, [CodeExplanationAndID(optimized_code, "Explanation", "2")])
    assert actual[0].cst_module.deep_equals(expected)

    original_code = '''class test_1():
    """original has it\n multi line"""
    pass
    def test_2(self):
        pass
    def test_3(self):
        """both have\n original"""
        pass'''
    original_module = cst.parse_module(original_code)
    optimized_code = libcst.parse_module(
        '''class test_1():
    pass
    def test_2(self):
        """retain this one v2\n multi line"""
        pass
    def test_3(self):
        """both have\n original v2"""
        pass'''
    )
    expected = libcst.parse_module(
        '''class test_1():
    """original has it\n multi line"""
    pass
    def test_2(self):
        """retain this one v2\n multi line"""
        pass
    def test_3(self):
        """both have\n original"""
        pass'''
    )
    actual = fix_missing_docstring(original_module, [CodeExplanationAndID(optimized_code, "Explanation", "2")])
    assert actual[0].cst_module.deep_equals(expected)

    original_code = '''class test_1:
    """both have it\n multi line"""
    pass
    def test_2(self):
        pass
    def test_3(self):
        """both have\n original"""
        pass'''
    original_module = cst.parse_module(original_code)
    optimized_code = libcst.parse_module(
        '''class test_1:
    """both have it\n multi line v2"""
    pass
    def test_2(self):
        """retain this one v2\n multi line"""
        pass
    def test_3(self):
        """both have\n original v2"""
        pass'''
    )
    expected = libcst.parse_module(
        '''class test_1:
    """both have it\n multi line"""
    pass
    def test_2(self):
        """retain this one v2\n multi line"""
        pass
    def test_3(self):
        """both have\n original"""
        pass'''
    )
    actual = fix_missing_docstring(original_module, [CodeExplanationAndID(optimized_code, "Explanation", "2")])
    assert actual[0].cst_module.deep_equals(expected)


def test_dedup_and_sort_imports_deduplicates() -> None:
    original_code = """
import os
import sys


def foo():
    return os.path.join(sys.path[0], 'bar')
"""
    original_module = cst.parse_module(original_code)
    optimizations = CodeExplanationAndID(
        libcst.parse_module(
            """
import os
import sys
import os


def foo():
    return os.path.join(sys.path[0], 'bar')
"""
        ),
        "Removed duplicate imports",
        "1",
    )

    expected = CodeExplanationAndID(
        libcst.parse_module(
            """
import os
import sys


def foo():
    return os.path.join(sys.path[0], 'bar')
"""
        ),
        "Removed duplicate imports",
        "1",
    )

    actual = dedup_and_sort_imports(original_module, [optimizations])

    assert actual[0].cst_module.deep_equals(expected.cst_module)


def test_dedup_and_sort_imports_sorts_and_deduplicates() -> None:
    original_code = """
import os
import sys
import json
import os


def foo():
    return os.path.join(sys.path[0], 'bar')
"""
    original_module = cst.parse_module(original_code)
    optimizations = CodeExplanationAndID(
        libcst.parse_module(
            """
import sys
import os
import json
import os


def foo():
    return os.path.join(sys.path[0], 'bar')
"""
        ),
        "Sorted and removed duplicate imports",
        "1",
    )

    expected = CodeExplanationAndID(
        libcst.parse_module(
            """
import json
import os
import sys


def foo():
    return os.path.join(sys.path[0], 'bar')
"""
        ),
        "Sorted and removed duplicate imports",
        "1",
    )

    actual = dedup_and_sort_imports(original_module, [optimizations])

    assert actual[0].cst_module.deep_equals(expected.cst_module)


def test_ellipsis_filter_removes_optimizations_with_ellipsis() -> None:
    original_code = """def foo():
    return True"""
    original_module = cst.parse_module(original_code)

    optimizations = [
        CodeExplanationAndID(
            libcst.parse_module(
                """def foo():
    return ..."""
            ),
            "Optimized foo",
            "1",
        ),
        CodeExplanationAndID(
            libcst.parse_module(
                """def foo():
        return False"""
            ),
            "Optimized foobar",
            "2",
        ),
    ]

    expected = [
        CodeExplanationAndID(
            libcst.parse_module(
                """def foo():
        return False"""
            ),
            "Optimized foobar",
            "2",
        )
    ]
    actual = filter_ellipsis_containing_code(original_module, optimizations)
    assert len(actual) == 1
    assert actual[0].cst_module.deep_equals(expected[0].cst_module)

    original_code2 = """def foo():
        return ..."""
    original_module2 = cst.parse_module(original_code2)
    actual = filter_ellipsis_containing_code(original_module2, optimizations)
    assert len(actual) == 2
    assert actual[0].cst_module.deep_equals(optimizations[0].cst_module)
    assert actual[1].cst_module.deep_equals(optimizations[1].cst_module)


# Doesn't work with isort, but maybe we don't care for now
#
# def test_dedup_and_sort_imports_sorts_multiple_blocks():
#     original_code = """
# import os
# import sys
#
# def foo():
#     return os.path.join(sys.path[0], 'bar')
#
# import json
# import os
# """
#     optimizations = [
#         (
#             """
# import sys
# import os
#
# def foo():
#     return os.path.join(sys.path[0], 'bar')
#
# import os
# import json
# """,
#             "Sorted imports and ensured they are not duplicated across multiple blocks",
#         ),
#     ]
#
#     expected = [
#         (
#             """
# import json
# import os
# import sys
#
# def foo():
#     return os.path.join(sys.path[0], 'bar')
# """,
#             "Sorted imports and ensured they are not duplicated across multiple blocks",
#         )
#     ]
#
#     actual = dedup_and_sort_imports(original_code, optimizations)
#
#     assert actual == expected
