from __future__ import annotations

from typing import List

import libcst as cst
# imports
import pytest  # used for our unit tests
from codeflash.code_utils.code_extractor import add_global_assignments


class ImportInserter(cst.CSTTransformer):
    def __init__(self, statements, last_import_line):
        self.statements = statements
        self.last_import_line = last_import_line

    def leave_Module(self, original_node, updated_node):
        new_body = list(updated_node.body)
        for statement in self.statements:
            new_body.insert(self.last_import_line, statement)
        return updated_node.with_changes(body=new_body)

class GlobalAssignmentCollector(cst.CSTVisitor):
    def __init__(self):
        self.assignments = []
        self.assignment_order = []

    def visit_Assign(self, node: cst.Assign) -> None:
        self.assignments.append(node)
        self.assignment_order.append(node)

class GlobalAssignmentTransformer(cst.CSTTransformer):
    def __init__(self, assignments, assignment_order):
        self.assignments = assignments
        self.assignment_order = assignment_order

    def leave_Module(self, original_node, updated_node):
        new_body = list(updated_node.body)
        for assignment in self.assignments:
            new_body.append(assignment)
        return updated_node.with_changes(body=new_body)

# unit tests

def test_simple_global_assignment():
    # Test with a simple global assignment
    src_code = "x = 1"
    dst_code = ""
    expected = "x = 1"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_multiple_global_assignments():
    # Test with multiple global assignments
    src_code = "a = 1\nb = 2"
    dst_code = ""
    expected = "a = 1\nb = 2"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_no_global_assignments_in_source():
    # Test with no global assignments in source
    src_code = ""
    dst_code = "import os"
    expected = "import os"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_no_import_statements_in_target():
    # Test with no import statements in target
    src_code = "x = 1"
    dst_code = "def foo(): pass"
    expected = "def foo(): pass\nx = 1"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_large_number_of_global_assignments():
    # Test with a large number of global assignments
    src_code = "\n".join(f"x{i} = {i}" for i in range(1000))
    dst_code = ""
    expected = "\n".join(f"x{i} = {i}" for i in range(1000))
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_conflicting_assignments():
    # Test with conflicting assignments
    src_code = "x = 1"
    dst_code = "x = 2"
    expected = "x = 2\nx = 1"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_comments_and_whitespace():
    # Test handling of comments and whitespace
    src_code = "x = 1  # comment"
    dst_code = "import os\n\n"
    expected = "import os\n\nx = 1  # comment"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_invalid_syntax():
    # Test handling of invalid syntax
    src_code = "x = "
    dst_code = "import os"
    with pytest.raises(Exception):
        add_global_assignments(src_code, dst_code)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from __future__ import annotations

from typing import List

import libcst as cst
# imports
import pytest  # used for our unit tests
from codeflash.code_utils.code_extractor import add_global_assignments


class ImportInserter(cst.CSTTransformer):
    def __init__(self, statements, last_import_line):
        self.statements = statements
        self.last_import_line = last_import_line

    def leave_Module(self, original_node, updated_node):
        body = list(updated_node.body)
        insertion_point = self.last_import_line + 1
        body[insertion_point:insertion_point] = self.statements
        return updated_node.with_changes(body=body)


class GlobalAssignmentCollector(cst.CSTVisitor):
    def __init__(self):
        self.assignments = []
        self.assignment_order = []

    def visit_Assign(self, node: cst.Assign):
        self.assignments.append(node)
        self.assignment_order.append(node.start.line)


class GlobalAssignmentTransformer(cst.CSTTransformer):
    def __init__(self, assignments, assignment_order):
        self.assignments = assignments
        self.assignment_order = assignment_order

    def leave_Module(self, original_node, updated_node):
        body = list(updated_node.body)
        for assignment, order in zip(self.assignments, self.assignment_order):
            body.insert(order, assignment)
        return updated_node.with_changes(body=body)


# unit tests

def test_simple_global_assignment():
    # Test with simple global assignment
    src_code = "x = 10"
    dst_code = ""
    expected = "x = 10"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_multiple_global_assignments():
    # Test with multiple global assignments
    src_code = "x = 10\ny = 20"
    dst_code = "z = 30"
    expected = "z = 30\nx = 10\ny = 20"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_empty_source_code():
    # Test with empty source code
    src_code = ""
    dst_code = "z = 30"
    expected = "z = 30"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_empty_destination_code():
    # Test with empty destination code
    src_code = "x = 10"
    dst_code = ""
    expected = "x = 10"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_no_global_assignments():
    # Test with no global assignments in source
    src_code = "import os"
    dst_code = "z = 30"
    expected = "z = 30"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_only_import_statements():
    # Test with only import statements in source
    src_code = "import os"
    dst_code = "x = 10"
    expected = "x = 10"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_nested_functions():
    # Test with nested functions in source
    src_code = "def outer():\n    x = 10\ndef inner():\n    pass"
    dst_code = "y = 20"
    expected = "y = 20\ndef outer():\n    x = 10\ndef inner():\n    pass"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_conditional_global_assignments():
    # Test with conditional global assignments
    src_code = "if True:\n    x = 10"
    dst_code = "y = 20"
    expected = "y = 20\nif True:\n    x = 10"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_global_assignments_with_comments():
    # Test with global assignments and comments
    src_code = "# This is a comment\nx = 10"
    dst_code = "y = 20"
    expected = "y = 20\n# This is a comment\nx = 10"
    codeflash_output = add_global_assignments(src_code, dst_code)


def test_line_continuations():
    # Test with line continuations
    src_code = "x = (10 +\n     20)"
    dst_code = "y = 30"
    expected = "y = 30\nx = (10 +\n     20)"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_large_source_code():
    # Test with large source code
    src_code = "\n".join(f"x{i} = {i}" for i in range(100))
    dst_code = "z = 100"
    expected = "z = 100\n" + "\n".join(f"x{i} = {i}" for i in range(100))
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_large_destination_code():
    # Test with large destination code
    src_code = "x = 10"
    dst_code = "\n".join(f"y{i} = {i}" for i in range(100))
    expected = "\n".join(f"y{i} = {i}" for i in range(100)) + "\nx = 10"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_performance_with_large_files():
    # Test performance with large files
    src_code = "\n".join(f"x{i} = {i}" for i in range(500))
    dst_code = "\n".join(f"y{i} = {i}" for i in range(500))
    codeflash_output = add_global_assignments(src_code, dst_code); result = codeflash_output

def test_malformed_source_code():
    # Test with malformed source code
    src_code = "x = 10\ny = "
    dst_code = "z = 30"
    with pytest.raises(cst.ParserSyntaxError):
        add_global_assignments(src_code, dst_code)

def test_malformed_destination_code():
    # Test with malformed destination code
    src_code = "x = 10"
    dst_code = "z = 30\ny = "
    with pytest.raises(cst.ParserSyntaxError):
        add_global_assignments(src_code, dst_code)

def test_imports_with_aliases():
    # Test with imports using aliases
    src_code = "import os as operating_system"
    dst_code = "x = 10"
    expected = "x = 10\nimport os as operating_system"
    codeflash_output = add_global_assignments(src_code, dst_code)

def test_circular_imports():
    # Test with potential circular imports
    src_code = "import module_a"
    dst_code = "import module_b"
    expected = "import module_b\nimport module_a"
    codeflash_output = add_global_assignments(src_code, dst_code)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from codeflash.code_utils.code_extractor import add_global_assignments

def test_add_global_assignments():
    add_global_assignments('', '')