from __future__ import annotations

import ast
import textwrap
from typing import TYPE_CHECKING

import libcst as cst
import pytest

if TYPE_CHECKING:
    from pathlib import Path

from codeflash_python._model import FunctionParent, FunctionToOptimize
from codeflash_python.analysis._code_utils import find_preexisting_objects
from codeflash_python.codegen._replacement import (
    DottedImportCollector,
    FutureAliasedImportTransformer,
    GlobalAssignmentCollector,
    GlobalAssignmentTransformer,
    GlobalFunctionCollector,
    GlobalFunctionTransformer,
    GlobalStatementCollector,
    GlobalStatementTransformer,
    add_global_assignments,
    add_needed_imports_from_module,
    collect_referenced_names,
    delete_future_aliased_imports,
    extract_global_statements,
    find_insertion_index_after_imports,
    gather_source_imports,
    is_zero_diff,
    normalize_code,
    normalize_node,
    replace_function_source,
    replace_functions_and_add_imports,
    replace_functions_in_file,
    resolve_star_import,
)


class TestReplaceFunctionSource:
    """Tests for replace_function_source."""

    def test_replace_top_level_function(self) -> None:
        """A top-level function body is replaced."""
        source = textwrap.dedent("""\
            import os

            def greet(name):
                return f"hello {name}"

            x = 1
        """)
        new_source = textwrap.dedent("""\
            def greet(name):
                return f"hi {name}"
        """)
        fn = FunctionToOptimize(
            function_name="greet",
            file_path="/dev/null",
            starting_line=3,
            ending_line=4,
        )

        result = replace_function_source(source, fn, new_source)

        assert 'return f"hi {name}"' in result
        assert 'return f"hello {name}"' not in result
        assert "import os" in result
        assert "x = 1" in result

    def test_replace_method(self) -> None:
        """A class method body is replaced while preserving the class."""
        source = textwrap.dedent("""\
            class Formatter:
                def bold(self, text):
                    return f"**{text}**"

                def italic(self, text):
                    return f"*{text}*"
        """)
        new_source = textwrap.dedent("""\
            class Formatter:
                def bold(self, text):
                    return "<b>" + text + "</b>"
        """)
        fn = FunctionToOptimize(
            function_name="bold",
            file_path="/dev/null",
            parents=(FunctionParent("Formatter", "ClassDef"),),
            starting_line=2,
            ending_line=3,
            is_method=True,
        )

        result = replace_function_source(source, fn, new_source)

        assert '"<b>" + text + "</b>"' in result
        assert "**{text}**" not in result
        # italic should be untouched
        assert "*{text}*" in result

    def test_preserves_surrounding_code(self) -> None:
        """Code before and after the function is not altered."""
        source = textwrap.dedent("""\
            CONSTANT = 42

            def compute(x):
                return x + 1

            def other(y):
                return y * 2
        """)
        new_source = textwrap.dedent("""\
            def compute(x):
                return x + 2
        """)
        fn = FunctionToOptimize(
            function_name="compute",
            file_path="/dev/null",
            starting_line=3,
            ending_line=4,
        )

        result = replace_function_source(source, fn, new_source)

        assert "CONSTANT = 42" in result
        assert "return x + 2" in result
        assert "return y * 2" in result

    def test_replaces_decorators(self) -> None:
        """Decorators from the new source replace the originals."""
        source = textwrap.dedent("""\
            @cache
            def fib(n):
                if n <= 1:
                    return n
                return fib(n - 1) + fib(n - 2)
        """)
        new_source = textwrap.dedent("""\
            @lru_cache(maxsize=128)
            def fib(n):
                if n <= 1:
                    return n
                return fib(n - 1) + fib(n - 2)
        """)
        fn = FunctionToOptimize(
            function_name="fib",
            file_path="/dev/null",
            starting_line=1,
            ending_line=5,
        )

        result = replace_function_source(source, fn, new_source)

        assert "@lru_cache(maxsize=128)" in result
        assert "@cache\n" not in result

    def test_function_not_in_new_source_raises(self) -> None:
        """ValueError raised when new_source doesn't contain the function."""
        fn = FunctionToOptimize(
            function_name="missing",
            file_path="/dev/null",
            starting_line=1,
            ending_line=2,
        )

        with pytest.raises(
            ValueError,
            match="not found in new_source",
        ):
            replace_function_source(
                "def missing(): return 1\n",
                fn,
                "def other(): return 2\n",
            )

    def test_round_trip_compiles(self) -> None:
        """Replaced code is valid Python."""
        source = textwrap.dedent("""\
            class Math:
                @staticmethod
                def add(a, b):
                    return a + b
        """)
        new_source = textwrap.dedent("""\
            class Math:
                @staticmethod
                def add(a, b):
                    return b + a
        """)
        fn = FunctionToOptimize(
            function_name="add",
            file_path="/dev/null",
            parents=(FunctionParent("Math", "ClassDef"),),
            starting_line=2,
            ending_line=4,
            is_method=True,
        )

        result = replace_function_source(source, fn, new_source)

        compile(result, "<test>", "exec")
        assert "return b + a" in result


class TestNormalizeNode:
    """Tests for normalize_node."""

    def test_strips_docstring_from_module(self) -> None:
        """Module-level docstring is removed."""
        code = textwrap.dedent("""\
            \"\"\"Module docstring.\"\"\"
            x = 1
        """)
        tree = ast.parse(code)
        result = normalize_node(tree)
        unparsed = ast.unparse(result)
        assert "Module docstring" not in unparsed
        assert "x = 1" in unparsed

    def test_strips_docstring_from_function(self) -> None:
        """Function-level docstring is removed."""
        code = textwrap.dedent("""\
            def foo():
                \"\"\"Function docstring.\"\"\"
                return 1
        """)
        tree = ast.parse(code)
        result = normalize_node(tree)
        unparsed = ast.unparse(result)
        assert "Function docstring" not in unparsed
        assert "return 1" in unparsed

    def test_removes_import_statements(self) -> None:
        """Import statements are filtered out."""
        code = textwrap.dedent("""\
            import os
            x = 1
        """)
        tree = ast.parse(code)
        result = normalize_node(tree)
        unparsed = ast.unparse(result)
        assert "import os" not in unparsed
        assert "x = 1" in unparsed

    def test_removes_from_import_statements(self) -> None:
        """From-import statements are filtered out."""
        code = textwrap.dedent("""\
            from os.path import join
            x = 1
        """)
        tree = ast.parse(code)
        result = normalize_node(tree)
        unparsed = ast.unparse(result)
        assert "from os.path" not in unparsed
        assert "x = 1" in unparsed

    def test_preserves_regular_code(self) -> None:
        """Assignments and function calls are kept intact."""
        code = textwrap.dedent("""\
            x = 1
            print(x)
        """)
        tree = ast.parse(code)
        result = normalize_node(tree)
        unparsed = ast.unparse(result)
        assert "x = 1" in unparsed
        assert "print(x)" in unparsed


class TestNormalizeCode:
    """Tests for normalize_code."""

    def test_identical_code_normalizes_to_same_string(self) -> None:
        """Identical code returns the same normalized string."""
        code = "x = 1\ny = 2\n"
        assert normalize_code(code) == normalize_code(code)

    def test_differing_docstrings_normalize_to_same(self) -> None:
        """Code differing only by docstrings normalizes identically."""
        a = textwrap.dedent("""\
            def foo():
                \"\"\"Old docstring.\"\"\"
                return 1
        """)
        b = textwrap.dedent("""\
            def foo():
                \"\"\"New docstring.\"\"\"
                return 1
        """)
        assert normalize_code(a) == normalize_code(b)

    def test_differing_imports_normalize_to_same(self) -> None:
        """Code differing only by imports normalizes identically."""
        a = textwrap.dedent("""\
            import os
            x = 1
        """)
        b = textwrap.dedent("""\
            import sys
            x = 1
        """)
        assert normalize_code(a) == normalize_code(b)

    def test_different_logic_produces_different_normalized(self) -> None:
        """Code with different logic does not normalize to the same string."""
        a = "x = 1\n"
        b = "x = 2\n"
        assert normalize_code(a) != normalize_code(b)


class TestIsZeroDiff:
    """Tests for is_zero_diff."""

    def test_identical_code_is_zero_diff(self) -> None:
        """Identical code is detected as zero diff."""
        code = "x = 1\n"
        assert is_zero_diff(code, code) is True

    def test_same_logic_different_formatting(self) -> None:
        """Whitespace-only changes are zero diff."""
        a = "x=1\n"
        b = "x = 1\n"
        assert is_zero_diff(a, b) is True

    def test_different_docstrings_is_zero_diff(self) -> None:
        """Differing only in docstrings is zero diff."""
        a = textwrap.dedent("""\
            def foo():
                \"\"\"Old.\"\"\"
                return 1
        """)
        b = textwrap.dedent("""\
            def foo():
                \"\"\"New.\"\"\"
                return 1
        """)
        assert is_zero_diff(a, b) is True

    def test_different_imports_is_zero_diff(self) -> None:
        """Differing only in imports is zero diff."""
        a = "import os\nx = 1\n"
        b = "import sys\nx = 1\n"
        assert is_zero_diff(a, b) is True

    def test_different_code_is_not_zero_diff(self) -> None:
        """Actually different logic is not zero diff."""
        a = "x = 1\n"
        b = "x = 2\n"
        assert is_zero_diff(a, b) is False


class TestFindPreexistingObjects:
    """Tests for find_preexisting_objects."""

    def test_finds_top_level_function(self) -> None:
        """A top-level function is returned as (name, ())."""
        code = "def func_name(): pass\n"
        result = find_preexisting_objects(code)
        assert ("func_name", ()) in result

    def test_finds_top_level_async_function(self) -> None:
        """A top-level async function is returned as (name, ())."""
        code = "async def async_func(): pass\n"
        result = find_preexisting_objects(code)
        assert ("async_func", ()) in result

    def test_finds_class(self) -> None:
        """A top-level class is returned as (name, ())."""
        code = "class ClassName: pass\n"
        result = find_preexisting_objects(code)
        assert ("ClassName", ()) in result

    def test_finds_class_method(self) -> None:
        """A class method is returned with FunctionParent tuple."""
        code = textwrap.dedent("""\
            class ClassName:
                def method(self):
                    pass
        """)
        result = find_preexisting_objects(code)
        expected = ("method", (FunctionParent("ClassName", "ClassDef"),))
        assert expected in result

    def test_empty_source_returns_empty_set(self) -> None:
        """Empty source code returns an empty set."""
        result = find_preexisting_objects("")
        assert set() == result

    def test_syntax_error_returns_empty_set(self) -> None:
        """Syntax error in source returns empty set without raising."""
        result = find_preexisting_objects("def ??? invalid syntax")
        assert set() == result


class TestFindInsertionIndexAfterImports:
    """Tests for find_insertion_index_after_imports."""

    def test_returns_zero_for_no_imports(self) -> None:
        """Module with no imports returns index 0."""
        module = cst.parse_module("x = 1\n")
        assert 0 == find_insertion_index_after_imports(module)

    def test_returns_index_after_single_import(self) -> None:
        """Index is 1 when there is one import at position 0."""
        module = cst.parse_module("import os\nx = 1\n")
        assert 1 == find_insertion_index_after_imports(module)

    def test_returns_index_after_last_import(self) -> None:
        """Index points past all consecutive imports."""
        code = textwrap.dedent("""\
            import os
            import sys
            from pathlib import Path
            x = 1
        """)
        module = cst.parse_module(code)
        assert 3 == find_insertion_index_after_imports(module)

    def test_stops_at_first_function(self) -> None:
        """Imports after a function definition are not counted."""
        code = textwrap.dedent("""\
            import os

            def foo():
                pass

            import sys
        """)
        module = cst.parse_module(code)
        assert 1 == find_insertion_index_after_imports(module)

    def test_stops_at_first_class(self) -> None:
        """Imports after a class definition are not counted."""
        code = textwrap.dedent("""\
            import os

            class Foo:
                pass

            import sys
        """)
        module = cst.parse_module(code)
        assert 1 == find_insertion_index_after_imports(module)

    def test_handles_conditional_import_block(self) -> None:
        """An if block containing only imports counts as an import."""
        code = textwrap.dedent("""\
            import os
            if TYPE_CHECKING:
                import sys
            x = 1
        """)
        module = cst.parse_module(code)
        assert 2 == find_insertion_index_after_imports(module)

    def test_empty_module(self) -> None:
        """Empty module returns index 0."""
        module = cst.parse_module("")
        assert 0 == find_insertion_index_after_imports(module)


class TestCollectReferencedNames:
    """Tests for collect_referenced_names."""

    def test_finds_names_in_simple_expression(self) -> None:
        """Name nodes in a binary expression are collected."""
        expr = cst.parse_expression("x + y")
        assert {"x", "y"} == collect_referenced_names(expr)

    def test_returns_empty_for_literal(self) -> None:
        """A pure integer literal has no Name nodes."""
        expr = cst.parse_expression("42")
        assert set() == collect_referenced_names(expr)

    def test_finds_names_in_nested_expression(self) -> None:
        """Name nodes in nested calls and subscripts are found."""
        expr = cst.parse_expression("foo(bar[baz])")
        assert {"foo", "bar", "baz"} == collect_referenced_names(expr)

    def test_finds_names_in_attribute_access(self) -> None:
        """The root Name in an attribute chain is collected."""
        expr = cst.parse_expression("os.path.join")
        names = collect_referenced_names(expr)
        assert "os" in names
        assert "path" in names
        assert "join" in names

    def test_string_literal_has_no_names(self) -> None:
        """A string literal contains no Name nodes."""
        expr = cst.parse_expression('"hello"')
        assert set() == collect_referenced_names(expr)


class TestGlobalFunctionCollector:
    """Tests for GlobalFunctionCollector."""

    def test_collects_module_level_functions(self) -> None:
        """Module-level functions are collected with correct order."""
        code = textwrap.dedent("""\
            def foo():
                pass

            def bar():
                pass
        """)
        module = cst.parse_module(code)
        collector = GlobalFunctionCollector()
        module.visit(collector)
        assert ["foo", "bar"] == collector.function_order
        assert "foo" in collector.functions
        assert "bar" in collector.functions

    def test_skips_class_bodies(self) -> None:
        """Methods inside classes are not collected."""
        code = textwrap.dedent("""\
            class MyClass:
                def method(self):
                    pass

            def standalone():
                pass
        """)
        module = cst.parse_module(code)
        collector = GlobalFunctionCollector()
        module.visit(collector)
        assert ["standalone"] == collector.function_order
        assert "method" not in collector.functions

    def test_skips_nested_functions(self) -> None:
        """Functions nested inside other functions are not collected."""
        code = textwrap.dedent("""\
            def outer():
                def inner():
                    pass
        """)
        module = cst.parse_module(code)
        collector = GlobalFunctionCollector()
        module.visit(collector)
        assert ["outer"] == collector.function_order
        assert "inner" not in collector.functions

    def test_empty_module(self) -> None:
        """Empty module yields no functions."""
        module = cst.parse_module("")
        collector = GlobalFunctionCollector()
        module.visit(collector)
        assert [] == collector.function_order
        assert {} == collector.functions

    def test_deduplicates_redefined_functions(self) -> None:
        """A function defined twice appears once in order list."""
        code = textwrap.dedent("""\
            def foo():
                return 1

            def foo():
                return 2
        """)
        module = cst.parse_module(code)
        collector = GlobalFunctionCollector()
        module.visit(collector)
        assert ["foo"] == collector.function_order
        # The second definition overwrites the first
        assert "return 2" in module.code


class TestGlobalFunctionTransformer:
    """Tests for GlobalFunctionTransformer."""

    def test_replaces_existing_function(self) -> None:
        """An existing function is replaced with the new version."""
        dst_code = textwrap.dedent("""\
            def foo():
                return 1
        """)
        new_code = textwrap.dedent("""\
            def foo():
                return 2
        """)
        dst_module = cst.parse_module(dst_code)
        new_module = cst.parse_module(new_code)
        new_func = new_module.body[0]
        assert isinstance(new_func, cst.FunctionDef)

        transformer = GlobalFunctionTransformer(
            {"foo": new_func},
            ["foo"],
        )
        result = dst_module.visit(transformer)
        assert "return 2" in result.code
        assert "return 1" not in result.code

    def test_appends_new_function(self) -> None:
        """A function not in the module is appended after definitions."""
        dst_code = textwrap.dedent("""\
            import os

            def existing():
                pass
        """)
        new_code = textwrap.dedent("""\
            def brand_new():
                return 42
        """)
        dst_module = cst.parse_module(dst_code)
        new_module = cst.parse_module(new_code)
        new_func = new_module.body[0]
        assert isinstance(new_func, cst.FunctionDef)

        transformer = GlobalFunctionTransformer(
            {"brand_new": new_func},
            ["brand_new"],
        )
        result = dst_module.visit(transformer)
        assert "def brand_new" in result.code
        assert "return 42" in result.code
        # The existing function is still there
        assert "def existing" in result.code

    def test_replaces_one_and_appends_another(self) -> None:
        """One function is replaced and another is appended."""
        dst_code = textwrap.dedent("""\
            def foo():
                return 1
        """)
        src_code = textwrap.dedent("""\
            def foo():
                return 99

            def bar():
                return 42
        """)
        dst_module = cst.parse_module(dst_code)
        src_module = cst.parse_module(src_code)
        funcs = {}
        order = []
        for stmt in src_module.body:
            if isinstance(stmt, cst.FunctionDef):
                funcs[stmt.name.value] = stmt
                order.append(stmt.name.value)

        transformer = GlobalFunctionTransformer(funcs, order)
        result = dst_module.visit(transformer)
        assert "return 99" in result.code
        assert "return 1" not in result.code
        assert "def bar" in result.code

    def test_skips_class_methods(self) -> None:
        """Methods inside classes are not replaced."""
        dst_code = textwrap.dedent("""\
            class MyClass:
                def foo(self):
                    return 1
        """)
        new_code = textwrap.dedent("""\
            def foo():
                return 2
        """)
        dst_module = cst.parse_module(dst_code)
        new_module = cst.parse_module(new_code)
        new_func = new_module.body[0]
        assert isinstance(new_func, cst.FunctionDef)

        transformer = GlobalFunctionTransformer(
            {"foo": new_func},
            ["foo"],
        )
        result = dst_module.visit(transformer)
        # The class method should NOT be replaced
        assert "return 1" in result.code
        # The new function should be appended at module level
        assert "return 2" in result.code

    def test_no_changes_when_empty(self) -> None:
        """No new functions means module is unchanged."""
        code = "x = 1\n"
        module = cst.parse_module(code)
        transformer = GlobalFunctionTransformer({}, [])
        result = module.visit(transformer)
        assert code == result.code


class TestGlobalAssignmentCollector:
    """Tests for GlobalAssignmentCollector."""

    def test_collects_simple_assignment(self) -> None:
        """A top-level simple assignment is collected."""
        code = "x = 1\n"
        module = cst.parse_module(code)
        collector = GlobalAssignmentCollector()
        module.visit(collector)
        assert ["x"] == collector.assignment_order
        assert "x" in collector.assignments

    def test_collects_annotated_assignment(self) -> None:
        """A top-level annotated assignment with value is collected."""
        code = "x: int = 1\n"
        module = cst.parse_module(code)
        collector = GlobalAssignmentCollector()
        module.visit(collector)
        assert ["x"] == collector.assignment_order
        assert "x" in collector.assignments

    def test_skips_annotated_assignment_without_value(self) -> None:
        """An annotated assignment without a value is not collected."""
        code = "x: int\n"
        module = cst.parse_module(code)
        collector = GlobalAssignmentCollector()
        module.visit(collector)
        assert [] == collector.assignment_order

    def test_tracks_order(self) -> None:
        """Multiple assignments preserve insertion order."""
        code = textwrap.dedent("""\
            a = 1
            b = 2
            c = 3
        """)
        module = cst.parse_module(code)
        collector = GlobalAssignmentCollector()
        module.visit(collector)
        assert ["a", "b", "c"] == collector.assignment_order

    def test_skips_inside_if_blocks(self) -> None:
        """Assignments inside if blocks are not collected."""
        code = textwrap.dedent("""\
            x = 1
            if True:
                y = 2
        """)
        module = cst.parse_module(code)
        collector = GlobalAssignmentCollector()
        module.visit(collector)
        assert ["x"] == collector.assignment_order
        assert "y" not in collector.assignments

    def test_skips_inside_functions(self) -> None:
        """Assignments inside function bodies are not collected."""
        code = textwrap.dedent("""\
            x = 1
            def foo():
                y = 2
        """)
        module = cst.parse_module(code)
        collector = GlobalAssignmentCollector()
        module.visit(collector)
        assert ["x"] == collector.assignment_order
        assert "y" not in collector.assignments

    def test_skips_inside_classes(self) -> None:
        """Assignments inside class bodies are not collected."""
        code = textwrap.dedent("""\
            x = 1
            class Foo:
                y = 2
        """)
        module = cst.parse_module(code)
        collector = GlobalAssignmentCollector()
        module.visit(collector)
        assert ["x"] == collector.assignment_order
        assert "y" not in collector.assignments

    def test_deduplicates_redefined_names(self) -> None:
        """A name assigned twice appears once in the order list."""
        code = textwrap.dedent("""\
            x = 1
            x = 2
        """)
        module = cst.parse_module(code)
        collector = GlobalAssignmentCollector()
        module.visit(collector)
        assert ["x"] == collector.assignment_order


class TestGlobalAssignmentTransformer:
    """Tests for GlobalAssignmentTransformer."""

    def test_replaces_existing_assignment(self) -> None:
        """An existing assignment is replaced with the new value."""
        dst_code = "x = 1\n"
        new_assign = cst.parse_statement("x = 99\n")
        assert isinstance(new_assign, cst.SimpleStatementLine)
        assign_node = new_assign.body[0]
        assert isinstance(assign_node, cst.Assign)

        module = cst.parse_module(dst_code)
        transformer = GlobalAssignmentTransformer(
            {"x": assign_node},
            ["x"],
        )
        result = module.visit(transformer)
        assert "x = 99" in result.code
        assert "x = 1" not in result.code

    def test_adds_new_assignment_after_imports(self) -> None:
        """A new assignment is placed after imports."""
        dst_code = textwrap.dedent("""\
            import os

            def foo():
                pass
        """)
        new_assign = cst.parse_statement("NEW_VAR = 42\n")
        assert isinstance(new_assign, cst.SimpleStatementLine)
        assign_node = new_assign.body[0]
        assert isinstance(assign_node, cst.Assign)

        module = cst.parse_module(dst_code)
        transformer = GlobalAssignmentTransformer(
            {"NEW_VAR": assign_node},
            ["NEW_VAR"],
        )
        result = module.visit(transformer)
        assert "NEW_VAR = 42" in result.code

    def test_adds_referencing_assignment_after_defs(self) -> None:
        """Assignment referencing a module-level name goes after defs."""
        dst_code = textwrap.dedent("""\
            import os

            def compute():
                return 1
        """)
        # This assignment references "compute", a module-level name
        new_assign = cst.parse_statement("result = compute()\n")
        assert isinstance(new_assign, cst.SimpleStatementLine)
        assign_node = new_assign.body[0]
        assert isinstance(assign_node, cst.Assign)

        module = cst.parse_module(dst_code)
        transformer = GlobalAssignmentTransformer(
            {"result": assign_node},
            ["result"],
        )
        result = module.visit(transformer)
        code = result.code
        assert "result = compute()" in code
        # result should appear after the function definition
        func_pos = code.index("def compute")
        result_pos = code.index("result = compute()")
        assert result_pos > func_pos

    def test_does_not_replace_inside_if_blocks(self) -> None:
        """Assignments inside if blocks are left unchanged."""
        dst_code = textwrap.dedent("""\
            x = 1
            if True:
                x = 2
        """)
        new_assign = cst.parse_statement("x = 99\n")
        assert isinstance(new_assign, cst.SimpleStatementLine)
        assign_node = new_assign.body[0]
        assert isinstance(assign_node, cst.Assign)

        module = cst.parse_module(dst_code)
        transformer = GlobalAssignmentTransformer(
            {"x": assign_node},
            ["x"],
        )
        result = module.visit(transformer)
        code = result.code
        # Top-level x should be replaced
        assert "x = 99" in code
        # The x inside if should still be 2
        assert "x = 2" in code

    def test_no_changes_when_empty(self) -> None:
        """No new assignments means module is unchanged."""
        code = "x = 1\n"
        module = cst.parse_module(code)
        transformer = GlobalAssignmentTransformer({}, [])
        result = module.visit(transformer)
        assert code == result.code


class TestGlobalStatementCollector:
    """Tests for GlobalStatementCollector."""

    def test_collects_expression_statement(self) -> None:
        """A module-level expression like print() is collected."""
        code = textwrap.dedent("""\
            import os
            x = 1
            print("hello")
        """)
        module = cst.parse_module(code)
        collector = GlobalStatementCollector()
        module.visit(collector)
        assert 1 == len(collector.global_statements)

    def test_skips_imports(self) -> None:
        """Import statements are not collected."""
        code = textwrap.dedent("""\
            import os
            from sys import path
        """)
        module = cst.parse_module(code)
        collector = GlobalStatementCollector()
        module.visit(collector)
        assert [] == collector.global_statements

    def test_skips_assignments(self) -> None:
        """Assignment statements are not collected."""
        code = "x = 1\ny: int = 2\n"
        module = cst.parse_module(code)
        collector = GlobalStatementCollector()
        module.visit(collector)
        assert [] == collector.global_statements

    def test_skips_functions_and_classes(self) -> None:
        """Function and class definitions are not collected."""
        code = textwrap.dedent("""\
            def foo():
                print("inside")

            class Bar:
                pass
        """)
        module = cst.parse_module(code)
        collector = GlobalStatementCollector()
        module.visit(collector)
        assert [] == collector.global_statements

    def test_collects_multiple_statements(self) -> None:
        """Multiple non-import, non-assignment statements are collected."""
        code = textwrap.dedent("""\
            print("first")
            print("second")
        """)
        module = cst.parse_module(code)
        collector = GlobalStatementCollector()
        module.visit(collector)
        assert 2 == len(collector.global_statements)

    def test_does_not_collect_statements_inside_functions(self) -> None:
        """Statements inside function bodies are not collected."""
        code = textwrap.dedent("""\
            def foo():
                print("inside")

            print("outside")
        """)
        module = cst.parse_module(code)
        collector = GlobalStatementCollector()
        module.visit(collector)
        assert 1 == len(collector.global_statements)


class TestGlobalStatementTransformer:
    """Tests for GlobalStatementTransformer."""

    def test_appends_statements_at_end(self) -> None:
        """Collected statements are appended after all existing code."""
        dst_code = textwrap.dedent("""\
            import os

            def foo():
                pass
        """)
        stmt = cst.parse_statement('print("added")\n')
        assert isinstance(stmt, cst.SimpleStatementLine)

        module = cst.parse_module(dst_code)
        transformer = GlobalStatementTransformer([stmt])
        result = module.visit(transformer)
        code = result.code
        assert 'print("added")' in code
        # The added statement should be after the function
        func_pos = code.index("def foo")
        added_pos = code.index('print("added")')
        assert added_pos > func_pos

    def test_no_changes_when_empty_list(self) -> None:
        """An empty list of statements leaves the module unchanged."""
        code = "x = 1\n"
        module = cst.parse_module(code)
        transformer = GlobalStatementTransformer([])
        result = module.visit(transformer)
        assert code == result.code

    def test_appends_multiple_statements(self) -> None:
        """Multiple statements are all appended at the end."""
        dst_code = "x = 1\n"
        stmt1 = cst.parse_statement('print("a")\n')
        stmt2 = cst.parse_statement('print("b")\n')
        assert isinstance(stmt1, cst.SimpleStatementLine)
        assert isinstance(stmt2, cst.SimpleStatementLine)

        module = cst.parse_module(dst_code)
        transformer = GlobalStatementTransformer([stmt1, stmt2])
        result = module.visit(transformer)
        assert 'print("a")' in result.code
        assert 'print("b")' in result.code

    def test_preserves_existing_code(self) -> None:
        """Existing code in the module is not altered."""
        dst_code = textwrap.dedent("""\
            import os
            x = 1
        """)
        stmt = cst.parse_statement('print("new")\n')
        assert isinstance(stmt, cst.SimpleStatementLine)

        module = cst.parse_module(dst_code)
        transformer = GlobalStatementTransformer([stmt])
        result = module.visit(transformer)
        assert "import os" in result.code
        assert "x = 1" in result.code


class TestExtractGlobalStatements:
    """Tests for extract_global_statements."""

    def test_returns_module_and_statements(self) -> None:
        """Returns a parsed module and list of global statements."""
        code = textwrap.dedent("""\
            import os
            print("hello")
        """)
        module, stmts = extract_global_statements(code)
        assert isinstance(module, cst.Module)
        assert 1 == len(stmts)

    def test_empty_code_returns_no_statements(self) -> None:
        """Empty code returns empty list of statements."""
        module, stmts = extract_global_statements("")
        assert isinstance(module, cst.Module)
        assert [] == stmts

    def test_only_imports_returns_no_statements(self) -> None:
        """Code with only imports returns no global statements."""
        code = textwrap.dedent("""\
            import os
            from sys import path
        """)
        _, stmts = extract_global_statements(code)
        assert [] == stmts

    def test_mixed_code_extracts_correct_statements(self) -> None:
        """Only non-import, non-assignment, non-def statements are returned."""
        code = textwrap.dedent("""\
            import os
            x = 1
            print("hello")
            def foo():
                pass
            print("world")
        """)
        _, stmts = extract_global_statements(code)
        assert 2 == len(stmts)


class TestAddGlobalAssignments:
    """Tests for add_global_assignments."""

    def test_adds_new_function_from_src_to_dst(self) -> None:
        """A function in src but not in dst is added."""
        src = textwrap.dedent("""\
            def helper():
                return 42
        """)
        dst = textwrap.dedent("""\
            def main():
                pass
        """)
        result = add_global_assignments(src, dst)
        assert "def helper" in result
        assert "return 42" in result
        assert "def main" in result

    def test_adds_new_assignment_from_src_to_dst(self) -> None:
        """An assignment in src is transferred to dst."""
        src = "CONSTANT = 42\n"
        dst = textwrap.dedent("""\
            def foo():
                pass
        """)
        result = add_global_assignments(src, dst)
        assert "CONSTANT = 42" in result
        assert "def foo" in result

    def test_returns_unchanged_when_nothing_to_add(self) -> None:
        """dst is unchanged when src has nothing to add."""
        src = "import os\n"
        dst = textwrap.dedent("""\
            import sys
            x = 1
        """)
        result = add_global_assignments(src, dst)
        assert dst == result

    def test_does_not_duplicate_existing_function(self) -> None:
        """A function already in dst is not added again."""
        src = textwrap.dedent("""\
            def shared():
                return 1
        """)
        dst = textwrap.dedent("""\
            def shared():
                return 1
        """)
        result = add_global_assignments(src, dst)
        assert 1 == result.count("def shared")

    def test_deduplicates_global_statements(self) -> None:
        """A global statement already in dst is not appended again."""
        src = textwrap.dedent("""\
            print("hello")
        """)
        dst = textwrap.dedent("""\
            print("hello")
        """)
        result = add_global_assignments(src, dst)
        assert 1 == result.count('print("hello")')

    def test_adds_function_and_assignment_together(self) -> None:
        """Both a new function and a new assignment are added."""
        src = textwrap.dedent("""\
            THRESHOLD = 10

            def compute(x):
                return x * 2
        """)
        dst = textwrap.dedent("""\
            import os
        """)
        result = add_global_assignments(src, dst)
        assert "THRESHOLD = 10" in result
        assert "def compute" in result
        assert "import os" in result

    def test_adds_global_statement(self) -> None:
        """A non-import, non-assignment statement from src is appended."""
        src = textwrap.dedent("""\
            print("setup complete")
        """)
        dst = textwrap.dedent("""\
            import os
        """)
        result = add_global_assignments(src, dst)
        assert 'print("setup complete")' in result

    def test_result_is_valid_python(self) -> None:
        """The output of add_global_assignments compiles without errors."""
        src = textwrap.dedent("""\
            CONSTANT = 42

            def helper(x):
                return x + CONSTANT

            print("init")
        """)
        dst = textwrap.dedent("""\
            import os

            def main():
                return os.getcwd()
        """)
        result = add_global_assignments(src, dst)
        compile(result, "<test>", "exec")


class TestReplaceFunctionsInFile:
    """Tests for replace_functions_in_file."""

    def test_replace_single_top_level_function(self) -> None:
        """A single top-level function is replaced in place."""
        source = textwrap.dedent("""\
            import os

            def greet(name):
                return f"hello {name}"

            x = 1
        """)
        optimized = textwrap.dedent("""\
            def greet(name):
                return f"hi {name}"
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["greet"],
            optimized,
            preexisting,
        )

        compile(result, "<test>", "exec")
        assert 'return f"hi {name}"' in result
        assert 'return f"hello {name}"' not in result
        assert "import os" in result
        assert "x = 1" in result

    def test_replace_class_method(self) -> None:
        """A class method is replaced using 'ClassName.method' format."""
        source = textwrap.dedent("""\
            class Formatter:
                def bold(self, text):
                    return f"**{text}**"

                def italic(self, text):
                    return f"*{text}*"
        """)
        optimized = textwrap.dedent("""\
            class Formatter:
                def bold(self, text):
                    return "<b>" + text + "</b>"
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["Formatter.bold"],
            optimized,
            preexisting,
        )

        compile(result, "<test>", "exec")
        assert '"<b>" + text + "</b>"' in result
        assert "**{text}**" not in result
        # italic should be untouched
        assert "*{text}*" in result

    def test_unsupported_nested_name_returns_unchanged(self) -> None:
        """Names with more than one dot return source unchanged."""
        source = textwrap.dedent("""\
            def foo():
                return 1
        """)
        optimized = textwrap.dedent("""\
            def foo():
                return 2
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["a.b.c"],
            optimized,
            preexisting,
        )

        assert source == result

    def test_insert_new_helper_function(self) -> None:
        """A helper function not in preexisting_objects is inserted."""
        source = textwrap.dedent("""\
            def compute(x):
                return x + 1
        """)
        optimized = textwrap.dedent("""\
            def helper(x):
                return x * 2

            def compute(x):
                return helper(x) + 1
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["compute"],
            optimized,
            preexisting,
        )

        compile(result, "<test>", "exec")
        assert "def helper" in result
        assert "return helper(x) + 1" in result

    def test_insert_new_helper_class(self) -> None:
        """A new class not in preexisting_objects is inserted."""
        source = textwrap.dedent("""\
            def process(x):
                return x + 1
        """)
        optimized = textwrap.dedent("""\
            class Cache:
                data = {}

            def process(x):
                return Cache.data.get(x, x + 1)
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["process"],
            optimized,
            preexisting,
        )

        compile(result, "<test>", "exec")
        assert "class Cache" in result
        assert "Cache.data.get" in result

    def test_replace_init_method(self) -> None:
        """An __init__ method is replaced when the class is preexisting."""
        source = textwrap.dedent("""\
            class Widget:
                def __init__(self):
                    self.value = 0

                def get_value(self):
                    return self.value
        """)
        optimized = textwrap.dedent("""\
            class Widget:
                def __init__(self):
                    self.value = 0
                    self.cached = True

                def get_value(self):
                    return self.value
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["Widget.get_value"],
            optimized,
            preexisting,
        )

        compile(result, "<test>", "exec")
        assert "self.cached = True" in result

    def test_add_new_method_to_existing_class(self) -> None:
        """A new method not in preexisting_objects is added to the class."""
        source = textwrap.dedent("""\
            class Calculator:
                def add(self, a, b):
                    return a + b
        """)
        optimized = textwrap.dedent("""\
            class Calculator:
                def add(self, a, b):
                    return self.validate(a) + self.validate(b)

                def validate(self, x):
                    return int(x)
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["Calculator.add"],
            optimized,
            preexisting,
        )

        compile(result, "<test>", "exec")
        assert "def validate" in result
        assert "self.validate(a)" in result

    def test_preserve_preexisting_objects(self) -> None:
        """Functions already in preexisting_objects are not inserted again."""
        source = textwrap.dedent("""\
            def existing_helper():
                return 42

            def target():
                return 1
        """)
        optimized = textwrap.dedent("""\
            def existing_helper():
                return 42

            def target():
                return existing_helper()
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["target"],
            optimized,
            preexisting,
        )

        compile(result, "<test>", "exec")
        assert "return existing_helper()" in result
        # existing_helper should appear exactly once (not duplicated)
        assert 2 == result.count("existing_helper")

    def test_multiple_function_replacement(self) -> None:
        """Multiple functions can be replaced in a single call."""
        source = textwrap.dedent("""\
            def foo():
                return 1

            def bar():
                return 2
        """)
        optimized = textwrap.dedent("""\
            def foo():
                return 10

            def bar():
                return 20
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["foo", "bar"],
            optimized,
            preexisting,
        )

        compile(result, "<test>", "exec")
        assert "return 10" in result
        assert "return 20" in result
        assert "return 1\n" not in result
        assert "return 2\n" not in result

    def test_empty_preexisting_objects(self) -> None:
        """An empty preexisting_objects set still allows basic replacement."""
        source = textwrap.dedent("""\
            def compute(x):
                return x + 1
        """)
        optimized = textwrap.dedent("""\
            def compute(x):
                return x + 2
        """)

        result = replace_functions_in_file(
            source,
            ["compute"],
            optimized,
            set(),
        )

        compile(result, "<test>", "exec")
        assert "return x + 2" in result
        assert "return x + 1" not in result

    def test_result_is_valid_python(self) -> None:
        """The output compiles as valid Python across all node types."""
        source = textwrap.dedent("""\
            import math

            class Geometry:
                def area(self, r):
                    return math.pi * r ** 2

            def circumference(r):
                return 2 * math.pi * r
        """)
        optimized = textwrap.dedent("""\
            PI_CACHED = 3.141592653589793

            class LookupTable:
                data = {}

            class Geometry:
                def area(self, r):
                    return PI_CACHED * r * r

                def perimeter(self, r):
                    return 2 * PI_CACHED * r

            def circumference(r):
                return 2 * PI_CACHED * r

            def helper():
                return PI_CACHED
        """)
        preexisting = find_preexisting_objects(source)

        result = replace_functions_in_file(
            source,
            ["Geometry.area", "circumference"],
            optimized,
            preexisting,
        )

        compile(result, "<test>", "exec")
        assert "PI_CACHED * r * r" in result
        assert "2 * PI_CACHED * r" in result
        assert "class LookupTable" in result
        assert "def helper" in result
        assert "def perimeter" in result


class TestDottedImportCollector:
    """Tests for DottedImportCollector."""

    def test_from_import(self) -> None:
        """``from X import Y`` becomes ``X.Y``."""
        code = "from pathlib import Path\n"
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "pathlib.Path" in collector.imports

    def test_plain_import(self) -> None:
        """``import os`` becomes ``os``."""
        code = "import os\n"
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "os" in collector.imports

    def test_import_with_alias(self) -> None:
        """``import numpy as np`` becomes ``numpy.np``."""
        code = "import numpy as np\n"
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "numpy.np" in collector.imports

    def test_from_import_with_alias(self) -> None:
        """``from os.path import join as pjoin`` becomes ``os.path.pjoin``."""
        code = "from os.path import join as pjoin\n"
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "os.path.pjoin" in collector.imports

    def test_skips_function_bodies(self) -> None:
        """Imports inside functions are not collected."""
        code = textwrap.dedent("""\
            import os

            def f():
                import sys
        """)
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "os" in collector.imports
        assert "sys" not in collector.imports

    def test_skips_class_bodies(self) -> None:
        """Imports inside classes are not collected."""
        code = textwrap.dedent("""\
            import os

            class C:
                import sys
        """)
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "os" in collector.imports
        assert "sys" not in collector.imports

    def test_collects_from_if_block(self) -> None:
        """Imports inside ``if`` blocks are collected."""
        code = textwrap.dedent("""\
            import sys
            if sys.version_info >= (3, 11):
                from typing import Self
        """)
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "sys" in collector.imports
        assert "typing.Self" in collector.imports

    def test_collects_from_try_block(self) -> None:
        """Imports inside ``try`` blocks are collected."""
        code = textwrap.dedent("""\
            try:
                from fast_lib import speed
            except ImportError:
                pass
        """)
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "fast_lib.speed" in collector.imports

    def test_star_import_skipped(self) -> None:
        """``from X import *`` does not add entries."""
        code = "from os import *\n"
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert len(collector.imports) == 0

    def test_multiple_names_from_single_import(self) -> None:
        """Multiple names from one ``from`` import."""
        code = "from os.path import join, exists\n"
        module = cst.parse_module(code)
        collector = DottedImportCollector()
        module.visit(collector)
        assert "os.path.join" in collector.imports
        assert "os.path.exists" in collector.imports


class TestFutureAliasedImportTransformer:
    """Tests for FutureAliasedImportTransformer."""

    def test_removes_aliased_future_import(self) -> None:
        """Aliased ``__future__`` import is removed."""
        code = "from __future__ import annotations as ann\n"
        result = delete_future_aliased_imports(code)
        assert "annotations" not in result

    def test_keeps_unaliased_future_import(self) -> None:
        """Unaliased ``__future__`` import is kept."""
        code = "from __future__ import annotations\n"
        result = delete_future_aliased_imports(code)
        assert "annotations" in result

    def test_partial_alias_removal(self) -> None:
        """Only the aliased name is stripped; unaliased stays."""
        code = "from __future__ import annotations, division as d\n"
        result = delete_future_aliased_imports(code)
        assert "annotations" in result
        assert "division" not in result

    def test_non_future_import_untouched(self) -> None:
        """Regular aliased imports are not affected."""
        code = "from os.path import join as pjoin\n"
        result = delete_future_aliased_imports(code)
        assert "pjoin" in result

    def test_transformer_directly(self) -> None:
        """Transformer can be used directly on a parsed module."""
        code = "from __future__ import annotations as a\n"
        module = cst.parse_module(code)
        result = module.visit(FutureAliasedImportTransformer())
        assert "annotations" not in result.code


class TestResolveStarImport:
    """Tests for resolve_star_import."""

    def test_resolves_from_all_list(self, tmp_path: Path) -> None:
        """Uses ``__all__`` when defined."""
        pkg = tmp_path / "mypkg"
        pkg.mkdir()
        (pkg / "utils.py").write_text(
            textwrap.dedent("""\
            __all__ = ["foo", "bar"]
            def foo(): ...
            def bar(): ...
            def _priv(): ...
        """)
        )
        result = resolve_star_import("mypkg.utils", tmp_path)
        assert result == {"foo", "bar"}

    def test_resolves_public_names_no_all(self, tmp_path: Path) -> None:
        """Falls back to public names when no ``__all__``."""
        pkg = tmp_path / "mypkg"
        pkg.mkdir()
        (pkg / "utils.py").write_text(
            "def foo(): ...\ndef _priv(): ...\nVAL = 1\n",
        )
        result = resolve_star_import("mypkg.utils", tmp_path)
        assert "foo" in result
        assert "VAL" in result
        assert "_priv" not in result

    def test_resolves_init_file(self, tmp_path: Path) -> None:
        """Resolves package init files."""
        pkg = tmp_path / "mypkg"
        pkg.mkdir()
        (pkg / "__init__.py").write_text(
            '__all__ = ["Thing"]\nclass Thing: ...\n',
        )
        result = resolve_star_import("mypkg", tmp_path)
        assert result == {"Thing"}

    def test_missing_module_returns_empty(self, tmp_path: Path) -> None:
        """Missing module returns empty set."""
        result = resolve_star_import("nonexistent.module", tmp_path)
        assert result == set()

    def test_includes_class_and_assign(self, tmp_path: Path) -> None:
        """Classes and top-level assignments are included."""
        pkg = tmp_path / "mypkg"
        pkg.mkdir()
        (pkg / "models.py").write_text(
            "class Foo: ...\nBAR = 1\n_PRIV = 2\n",
        )
        result = resolve_star_import("mypkg.models", tmp_path)
        assert "Foo" in result
        assert "BAR" in result
        assert "_PRIV" not in result


class TestGatherSourceImports:
    """Tests for gather_source_imports."""

    def test_gathers_from_source(self, tmp_path: Path) -> None:
        """Imports are gathered from source module code."""
        src = tmp_path / "src.py"
        src.write_text("import os\nfrom pathlib import Path\ndef f(): ...\n")
        result = gather_source_imports(src.read_text(), src, tmp_path)
        assert result is not None

    def test_returns_none_for_no_imports(self, tmp_path: Path) -> None:
        """Returns None when no imports exist."""
        src = tmp_path / "src.py"
        src.write_text("def f(): ...\n")
        result = gather_source_imports(src.read_text(), src, tmp_path)
        assert result is None

    def test_accepts_cst_module(self, tmp_path: Path) -> None:
        """Accepts a pre-parsed cst.Module."""
        src = tmp_path / "src.py"
        src.write_text("import os\ndef f(): ...\n")
        module = cst.parse_module(src.read_text())
        result = gather_source_imports(module, src, tmp_path)
        assert result is not None

    def test_strips_future_aliases(self, tmp_path: Path) -> None:
        """Aliased __future__ imports are cleaned."""
        src = tmp_path / "src.py"
        src.write_text(
            textwrap.dedent("""\
            from __future__ import annotations as ann
            import os
            def f(): ...
        """)
        )
        result = gather_source_imports(src.read_text(), src, tmp_path)
        assert result is not None


class TestAddNeededImportsFromModule:
    """Tests for add_needed_imports_from_module."""

    def test_adds_missing_import(self, tmp_path: Path) -> None:
        """A missing import is added to the destination."""
        src = tmp_path / "src.py"
        src.write_text("import os\ndef f():\n    return os.getcwd()\n")
        dst = tmp_path / "dst.py"
        dst.write_text("def g():\n    return os.getcwd()\n")
        result = add_needed_imports_from_module(
            src.read_text(),
            dst.read_text(),
            src,
            dst,
            tmp_path,
        )
        assert "import os" in result

    def test_does_not_duplicate_existing_import(
        self,
        tmp_path: Path,
    ) -> None:
        """An already-present import is not duplicated."""
        src = tmp_path / "src.py"
        src.write_text(
            "import os\ndef f():\n    return os.getcwd()\n",
        )
        dst = tmp_path / "dst.py"
        dst.write_text(
            "import os\ndef g():\n    return os.listdir()\n",
        )
        result = add_needed_imports_from_module(
            src.read_text(),
            dst.read_text(),
            src,
            dst,
            tmp_path,
        )
        assert result.count("import os") == 1

    def test_handles_from_imports(
        self,
        tmp_path: Path,
    ) -> None:
        """``from X import Y`` style imports are added."""
        src = tmp_path / "src.py"
        src.write_text(
            textwrap.dedent("""\
            from pathlib import Path
            def f():
                return Path('.')
        """),
        )
        dst = tmp_path / "dst.py"
        dst.write_text(
            "def g():\n    return Path('.')\n",
        )
        result = add_needed_imports_from_module(
            src.read_text(),
            dst.read_text(),
            src,
            dst,
            tmp_path,
        )
        assert "from pathlib import Path" in result

    def test_no_imports_returns_original(self, tmp_path: Path) -> None:
        """When source has no imports, destination is returned unchanged."""
        src = tmp_path / "src.py"
        src.write_text("def f(): ...\n")
        dst = tmp_path / "dst.py"
        dst_code = "def g(): ...\n"
        dst.write_text(dst_code)
        result = add_needed_imports_from_module(
            src.read_text(),
            dst_code,
            src,
            dst,
            tmp_path,
        )
        assert result == dst_code

    def test_pre_gathered_imports(self, tmp_path: Path) -> None:
        """Accepts pre-gathered imports to avoid re-parsing source."""
        src = tmp_path / "src.py"
        src.write_text("import os\ndef f(): ...\n")
        dst = tmp_path / "dst.py"
        dst.write_text("def g():\n    return os.getcwd()\n")
        gathered = gather_source_imports(src.read_text(), src, tmp_path)
        result = add_needed_imports_from_module(
            src.read_text(),
            dst.read_text(),
            src,
            dst,
            tmp_path,
            gathered_imports=gathered,
        )
        assert "import os" in result


class TestReplaceFunctionsAndAddImports:
    """Tests for replace_functions_and_add_imports."""

    def test_replaces_and_adds_imports(self, tmp_path: Path) -> None:
        """Functions are replaced and needed imports are added."""
        source = textwrap.dedent("""\
            def area(r):
                return 3.14 * r * r
        """)
        optimized = textwrap.dedent("""\
            import math

            def area(r):
                return math.pi * r * r
        """)
        module = tmp_path / "mod.py"
        module.write_text(source)
        result = replace_functions_and_add_imports(
            source,
            ["area"],
            optimized,
            module,
            set(),
            tmp_path,
        )
        assert "math.pi" in result
        assert "import math" in result
        compile(result, "<test>", "exec")

    def test_preserves_existing_imports(self, tmp_path: Path) -> None:
        """Existing imports in source are not removed."""
        source = textwrap.dedent("""\
            import os

            def f():
                return os.getcwd()
        """)
        optimized = textwrap.dedent("""\
            import os

            def f():
                return os.path.abspath(".")
        """)
        module = tmp_path / "mod.py"
        module.write_text(source)
        result = replace_functions_and_add_imports(
            source,
            ["f"],
            optimized,
            module,
            set(),
            tmp_path,
        )
        assert "import os" in result
        assert 'os.path.abspath(".")' in result
        compile(result, "<test>", "exec")
