from __future__ import annotations

import libcst as cst
import pytest

from codeflash_api.languages.python._cst_utils import (
    AnyEllipsisVisitor,
    DefinitionRemover,
    DepthTrackingMixin,
    ImportTrackingVisitor,
    InvalidEllipsisVisitor,
    any_ellipsis_in_cst,
    build_module_path,
    collect_imported_names_from_import,
    collect_imported_names_from_import_from,
    compare_unparsed_ast_to_source,
    ellipsis_in_cst_not_types,
    evaluate_expression,
    extract_import_info,
    extract_imports_from_import,
    extract_imports_from_import_from,
    file_path_to_module_path,
    find_init,
    get_base_class_name,
    get_dotted_name,
    has_decorator,
    make_number_node,
    parse_module_to_cst,
    unparse_parse_source,
)
from codeflash_api.languages.python._postprocess import (
    CSTAnnotationNameCollector,
    DocstringTransformer,
    DocstringVisitor,
    EllipsisContainingCodeVisitor,
    OptimizationCandidate,
    _strip_comments_from_code,
    add_future_annotations_import,
    clean_extraneous_comments,
    clean_extraneous_comments_pipeline,
    cleanup_explanations,
    dedup_and_sort_imports,
    deduplicate_optimizations,
    equality_check,
    extract_names_from_cst_annotation,
    filter_ellipsis_containing_code,
    fix_forward_references,
    fix_missing_docstring,
    has_future_annotations_import,
    optimizations_postprocessing_pipeline,
    safe_isort,
)
from codeflash_api.languages.python._validator import (
    parse_python_or_none,
    validate_python_syntax,
)


class TestDepthTrackingMixin:
    """Tests for DepthTrackingMixin."""

    def test_initial_state(self) -> None:
        """
        Starts at top level.
        """
        mixin = DepthTrackingMixin()
        assert mixin._is_top_level()

    def test_function_depth(self) -> None:
        """
        Entering a function increases depth.
        """
        mixin = DepthTrackingMixin()
        mixin._visit_function()
        assert not mixin._is_top_level()
        assert mixin._is_inside_function()
        mixin._leave_function()
        assert mixin._is_top_level()


class TestFilePathToModulePath:
    """Tests for file_path_to_module_path."""

    def test_unix_path(self) -> None:
        """
        Forward slashes become dots.
        """
        assert "path.to.module" == file_path_to_module_path(
            "path/to/module.py"
        )

    def test_windows_path(self) -> None:
        """
        Backslashes become dots.
        """
        assert "path.to.module" == file_path_to_module_path(
            "path\\to\\module.py"
        )


class TestGetDottedName:
    """Tests for get_dotted_name."""

    def test_simple_name(self) -> None:
        """
        Simple Name node returns the value.
        """
        node = cst.Name("foo")
        assert "foo" == get_dotted_name(node)

    def test_attribute(self) -> None:
        """
        Dotted attribute returns full path.
        """
        node = cst.Attribute(value=cst.Name("foo"), attr=cst.Name("bar"))
        assert "foo.bar" == get_dotted_name(node)

    def test_none(self) -> None:
        """
        None returns empty string.
        """
        assert "" == get_dotted_name(None)


class TestBuildModulePath:
    """Tests for build_module_path."""

    def test_single_part(self) -> None:
        """
        Single name becomes cst.Name.
        """
        result = build_module_path("foo")
        assert isinstance(result, cst.Name)
        assert "foo" == result.value

    def test_dotted(self) -> None:
        """
        Dotted path becomes cst.Attribute chain.
        """
        result = build_module_path("foo.bar.baz")
        assert isinstance(result, cst.Attribute)
        assert "baz" == result.attr.value


class TestEvaluateExpression:
    """Tests for evaluate_expression."""

    def test_integer(self) -> None:
        """
        Integer node evaluates to its value.
        """
        assert 42 == evaluate_expression(cst.Integer("42"))

    def test_hex(self) -> None:
        """
        Hex integer evaluates correctly.
        """
        assert 255 == evaluate_expression(cst.Integer("0xff"))

    def test_negative(self) -> None:
        """
        Unary minus evaluates to negative.
        """
        node = cst.UnaryOperation(
            operator=cst.Minus(),
            expression=cst.Integer("5"),
        )
        assert -5 == evaluate_expression(node)

    def test_binary_add(self) -> None:
        """
        Addition evaluates correctly.
        """
        node = cst.BinaryOperation(
            left=cst.Integer("3"),
            operator=cst.Add(),
            right=cst.Integer("4"),
        )
        assert 7 == evaluate_expression(node)

    def test_unevaluable(self) -> None:
        """
        Name node returns None.
        """
        assert evaluate_expression(cst.Name("x")) is None


class TestMakeNumberNode:
    """Tests for make_number_node."""

    def test_positive(self) -> None:
        """
        Positive value becomes Integer.
        """
        node = make_number_node(5)
        assert isinstance(node, cst.Integer)
        assert "5" == node.value

    def test_negative(self) -> None:
        """
        Negative value becomes UnaryOperation.
        """
        node = make_number_node(-3)
        assert isinstance(node, cst.UnaryOperation)


class TestParseModuleToCst:
    """Tests for parse_module_to_cst."""

    def test_caches(self) -> None:
        """
        Same input returns same object.
        """
        code = "x = 1"
        a = parse_module_to_cst(code)
        b = parse_module_to_cst(code)
        assert a is b


class TestDefinitionRemover:
    """Tests for DefinitionRemover."""

    def test_removes_function(self) -> None:
        """
        Top-level function is removed by name.
        """
        code = "def foo():\n    pass\ndef bar():\n    pass\n"
        module = cst.parse_module(code)
        remover = DefinitionRemover({"foo"})
        result = module.visit(remover)
        assert "foo" not in result.code
        assert "bar" in result.code
        assert "foo" in remover.removed_names

    def test_protected_name(self) -> None:
        """
        Protected name is not removed.
        """
        code = "def foo():\n    pass\n"
        module = cst.parse_module(code)
        remover = DefinitionRemover({"foo"}, protected_names={"foo"})
        result = module.visit(remover)
        assert "foo" in result.code


class TestImportTrackingVisitor:
    """Tests for ImportTrackingVisitor."""

    def test_tracks_imports(self) -> None:
        """
        Imported names are collected.
        """
        import ast

        tree = ast.parse("import os\nfrom sys import path")
        visitor = ImportTrackingVisitor()
        visitor.visit(tree)
        assert "os" in visitor.imported_names
        assert "path" in visitor.imported_names


class TestFindInit:
    """Tests for find_init."""

    def test_found(self) -> None:
        """
        __init__ in a class is found.
        """
        import ast

        tree = ast.parse("class Foo:\n    def __init__(self):\n        pass\n")
        assert find_init(tree)

    def test_not_found(self) -> None:
        """
        No __init__ returns False.
        """
        import ast

        tree = ast.parse("class Foo:\n    def bar(self):\n        pass\n")
        assert not find_init(tree)


class TestValidatePythonSyntax:
    """Tests for validate_python_syntax."""

    def test_valid(self) -> None:
        """
        Valid Python returns True.
        """
        assert validate_python_syntax("x = 1\n")

    def test_invalid(self) -> None:
        """
        Invalid Python returns False.
        """
        assert not validate_python_syntax("def (:\n")

    def test_empty_body(self) -> None:
        """
        Empty content (only comments) returns False.
        """
        assert not validate_python_syntax("# just a comment\n")


class TestParsePythonOrNone:
    """Tests for parse_python_or_none."""

    def test_valid(self) -> None:
        """
        Valid Python returns a Module.
        """
        result = parse_python_or_none("x = 1\n")
        assert result is not None
        assert isinstance(result, cst.Module)

    def test_invalid(self) -> None:
        """
        Invalid Python returns None.
        """
        assert parse_python_or_none("def (:\n") is None


class TestOptimizationCandidate:
    """Tests for OptimizationCandidate."""

    def test_frozen(self) -> None:
        """
        OptimizationCandidate is immutable.
        """
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 1"),
            explanation="faster",
            id="test",
        )
        with pytest.raises(AttributeError):
            c.id = "changed"


class TestSafeIsort:
    """Tests for safe_isort."""

    def test_sorts(self) -> None:
        """
        Imports are sorted.
        """
        code = "import sys\nimport os\n"
        result = safe_isort(code)
        assert result.index("os") < result.index("sys")

    def test_invalid_returns_original(self) -> None:
        """
        Invalid code returns unchanged.
        """
        code = "not valid python {{{{"
        assert code == safe_isort(code)


class TestDeduplicateOptimizations:
    """Tests for deduplicate_optimizations."""

    def test_removes_duplicates(self) -> None:
        """
        Candidates with same AST are deduplicated.
        """
        module = cst.parse_module("x  =  1\n")
        c1 = OptimizationCandidate(
            cst_module=cst.parse_module("x = 1\n"),
            explanation="a",
            id="1",
        )
        c2 = OptimizationCandidate(
            cst_module=cst.parse_module("x=1\n"),
            explanation="b",
            id="2",
        )
        result = deduplicate_optimizations(module, [c1, c2])
        assert 1 == len(result)


class TestEqualityCheck:
    """Tests for equality_check."""

    def test_filters_identical(self) -> None:
        """
        Candidate identical to original is filtered.
        """
        module = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 1\n"),
            explanation="same",
            id="1",
        )
        result = equality_check(module, [c])
        assert 0 == len(result)

    def test_keeps_different(self) -> None:
        """
        Different candidate is kept.
        """
        module = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 2\n"),
            explanation="different",
            id="1",
        )
        result = equality_check(module, [c])
        assert 1 == len(result)


class TestFilterEllipsis:
    """Tests for filter_ellipsis_containing_code."""

    def test_filters_introduced_ellipsis(self) -> None:
        """
        Candidate that introduces ellipsis is filtered.
        """
        module = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = ...\n"),
            explanation="bad",
            id="1",
        )
        result = filter_ellipsis_containing_code(module, [c])
        assert 0 == len(result)

    def test_keeps_when_original_has_ellipsis(self) -> None:
        """
        If original has ellipsis, candidates are not filtered.
        """
        module = cst.parse_module("x = ...\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = ...\n"),
            explanation="ok",
            id="1",
        )
        result = filter_ellipsis_containing_code(module, [c])
        assert 1 == len(result)


class TestFixMissingDocstring:
    """Tests for fix_missing_docstring."""

    def test_restores_docstring(self) -> None:
        """
        Removed docstring is restored.
        """
        original = cst.parse_module(
            'def foo():\n    """My docstring."""\n    pass\n'
        )
        optimized = cst.parse_module("def foo():\n    pass\n")
        c = OptimizationCandidate(
            cst_module=optimized,
            explanation="faster",
            id="1",
        )
        result = fix_missing_docstring(original, [c])
        assert "My docstring" in result[0].cst_module.code


class TestHasFutureAnnotationsImport:
    """Tests for has_future_annotations_import."""

    def test_present(self) -> None:
        """
        Module with the import returns True.
        """
        module = cst.parse_module("from __future__ import annotations\n")
        assert has_future_annotations_import(module)

    def test_absent(self) -> None:
        """
        Module without the import returns False.
        """
        module = cst.parse_module("x = 1\n")
        assert not has_future_annotations_import(module)


class TestAddFutureAnnotationsImport:
    """Tests for add_future_annotations_import."""

    def test_adds_when_needed(self) -> None:
        """
        Import is added when undefined annotations exist.
        """
        code = "def foo(x: MyType) -> None:\n    pass\n"
        module = cst.parse_module(code)
        result = add_future_annotations_import(module)
        assert "from __future__ import annotations" in result.code

    def test_skips_when_present(self) -> None:
        """
        No change when import already exists.
        """
        code = "from __future__ import annotations\ndef foo(x: MyType) -> None:\n    pass\n"
        module = cst.parse_module(code)
        result = add_future_annotations_import(module)
        assert result is module

    def test_skips_when_no_undefined(self) -> None:
        """
        No change when all annotations are defined.
        """
        code = "def foo(x: int) -> None:\n    pass\n"
        module = cst.parse_module(code)
        result = add_future_annotations_import(module)
        assert result is module


class TestCleanupExplanations:
    """Tests for cleanup_explanations."""

    def test_removes_code_block(self) -> None:
        """
        Markdown code blocks are removed from explanations.
        """
        module = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=module,
            explanation="```python\nx = 1\n```",
            id="1",
        )
        result = cleanup_explanations(module, [c])
        assert "```" not in result[0].explanation


class TestPostprocessingPipeline:
    """Tests for optimizations_postprocessing_pipeline."""

    def test_empty_candidates(self) -> None:
        """
        Empty list returns empty list.
        """
        module = cst.parse_module("x = 1\n")
        result = optimizations_postprocessing_pipeline(module, [])
        assert [] == result

    def test_filters_identical(self) -> None:
        """
        Identical candidate is filtered by the pipeline.
        """
        module = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 1\n"),
            explanation="same",
            id="1",
        )
        result = optimizations_postprocessing_pipeline(module, [c])
        assert 0 == len(result)

    def test_keeps_valid_optimization(self) -> None:
        """
        Valid different candidate passes through.
        """
        module = cst.parse_module("x = 1 + 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 2\n"),
            explanation="precomputed",
            id="1",
        )
        result = optimizations_postprocessing_pipeline(module, [c])
        assert 1 == len(result)


# ---------------------------------------------------------------------------
# _cst_utils.py — additional coverage
# ---------------------------------------------------------------------------


class TestDepthTrackingMixinClassDepth:
    """Tests for class depth tracking in DepthTrackingMixin."""

    def test_class_depth(self) -> None:
        """Entering a class increases class depth."""
        mixin = DepthTrackingMixin()
        assert not mixin._is_inside_class()
        mixin._visit_class()
        assert mixin._is_inside_class()
        assert not mixin._is_top_level_class()
        mixin._leave_class()
        assert mixin._is_top_level_class()

    def test_top_level_function_check(self) -> None:
        """_is_top_level_function returns False inside a function."""
        mixin = DepthTrackingMixin()
        assert mixin._is_top_level_function()
        mixin._visit_function()
        assert not mixin._is_top_level_function()
        mixin._leave_function()
        assert mixin._is_top_level_function()

    def test_top_level_class_with_function(self) -> None:
        """_is_top_level_class returns False inside a function."""
        mixin = DepthTrackingMixin()
        mixin._visit_function()
        assert not mixin._is_top_level_class()


class TestExtractImportInfo:
    """Tests for extract_import_info."""

    def test_simple_name(self) -> None:
        """Simple import alias extracts name correctly."""
        alias = cst.ImportAlias(name=cst.Name("os"))
        available, module, original = extract_import_info(alias)
        assert "os" == available
        assert "" == module
        assert "os" == original

    def test_dotted_name_without_module(self) -> None:
        """Dotted import without module returns first component as available."""
        alias = cst.ImportAlias(
            name=cst.Attribute(value=cst.Name("os"), attr=cst.Name("path"))
        )
        available, module, original = extract_import_info(alias)
        assert "os" == available
        assert "" == module
        assert "os.path" == original

    def test_with_asname(self) -> None:
        """Import alias with asname uses the alias."""
        alias = cst.ImportAlias(
            name=cst.Name("os"),
            asname=cst.AsName(
                whitespace_before_as=cst.SimpleWhitespace(" "),
                whitespace_after_as=cst.SimpleWhitespace(" "),
                name=cst.Name("operating_system"),
            ),
        )
        available, module, original = extract_import_info(alias)
        assert "operating_system" == available
        assert "os" == original

    def test_with_module_name(self) -> None:
        """When module_name is given, available_name is the original name."""
        alias = cst.ImportAlias(name=cst.Name("path"))
        available, module, original = extract_import_info(alias, "os")
        assert "path" == available
        assert "os" == module
        assert "path" == original


class TestExtractImportsFromImport:
    """Tests for extract_imports_from_import."""

    def test_simple_import(self) -> None:
        """Extracts names from a plain import statement."""
        node = cst.parse_statement("import os, sys\n")
        assert isinstance(node, cst.SimpleStatementLine)
        import_node = node.body[0]
        assert isinstance(import_node, cst.Import)
        result = extract_imports_from_import(import_node)
        assert "os" in result
        assert "sys" in result


class TestExtractImportsFromImportFrom:
    """Tests for extract_imports_from_import_from."""

    def test_from_import(self) -> None:
        """Extracts names from 'from X import Y'."""
        node = cst.parse_statement("from os import path, getcwd\n")
        assert isinstance(node, cst.SimpleStatementLine)
        import_from = node.body[0]
        assert isinstance(import_from, cst.ImportFrom)
        result = extract_imports_from_import_from(import_from)
        assert "path" in result
        assert "getcwd" in result
        assert ("os", "path") == result["path"]

    def test_import_star_returns_empty(self) -> None:
        """ImportStar returns empty dict."""
        node = cst.parse_statement("from os import *\n")
        assert isinstance(node, cst.SimpleStatementLine)
        import_from = node.body[0]
        assert isinstance(import_from, cst.ImportFrom)
        result = extract_imports_from_import_from(import_from)
        assert {} == result


class TestCollectImportedNames:
    """Tests for collect_imported_names_from_import and _from_import_from."""

    def test_collect_from_import(self) -> None:
        """Collects available names from import statement."""
        node = cst.parse_statement("import os, sys\n")
        import_node = node.body[0]
        assert isinstance(import_node, cst.Import)
        result = collect_imported_names_from_import(import_node)
        assert {"os", "sys"} == result

    def test_collect_from_import_from(self) -> None:
        """Collects available names from from-import statement."""
        node = cst.parse_statement("from os import path, getcwd\n")
        import_from = node.body[0]
        assert isinstance(import_from, cst.ImportFrom)
        result = collect_imported_names_from_import_from(import_from)
        assert {"path", "getcwd"} == result

    def test_collect_from_import_from_star(self) -> None:
        """ImportStar from-import returns empty set."""
        node = cst.parse_statement("from os import *\n")
        import_from = node.body[0]
        assert isinstance(import_from, cst.ImportFrom)
        result = collect_imported_names_from_import_from(import_from)
        assert set() == result


class TestDefinitionRemoverAdvanced:
    """Advanced tests for DefinitionRemover."""

    def test_removes_class(self) -> None:
        """Top-level class is removed by name."""
        code = "class Foo:\n    pass\nclass Bar:\n    pass\n"
        module = cst.parse_module(code)
        remover = DefinitionRemover({"Foo"})
        result = module.visit(remover)
        assert "Foo" not in result.code
        assert "Bar" in result.code
        assert "Foo" in remover.removed_names

    def test_removes_method_causes_class_removal(self) -> None:
        """Removing a qualified method removes the enclosing class."""
        code = "class Foo:\n    def bar(self):\n        pass\n"
        module = cst.parse_module(code)
        remover = DefinitionRemover({"Foo.bar"})
        result = module.visit(remover)
        assert "Foo" not in result.code
        assert "Foo" in remover.removed_names

    def test_nested_function_not_removed(self) -> None:
        """A function nested inside another is not removed as top-level."""
        code = "def outer():\n    def foo():\n        pass\n    return foo()\n"
        module = cst.parse_module(code)
        remover = DefinitionRemover({"foo"})
        result = module.visit(remover)
        assert "foo" in result.code

    def test_protected_class_not_removed(self) -> None:
        """Protected class is not removed."""
        code = "class Foo:\n    pass\n"
        module = cst.parse_module(code)
        remover = DefinitionRemover({"Foo"}, protected_names={"Foo"})
        result = module.visit(remover)
        assert "Foo" in result.code


class TestImportTrackingVisitorAdvanced:
    """Additional tests for ImportTrackingVisitor."""

    def test_tracks_import_with_alias(self) -> None:
        """Aliased imports use the alias name."""
        import ast

        tree = ast.parse("import os as operating_system")
        visitor = ImportTrackingVisitor()
        visitor.visit(tree)
        assert "operating_system" in visitor.imported_names
        assert "os" not in visitor.imported_names

    def test_tracks_from_import_with_alias(self) -> None:
        """Aliased from-imports use the alias name."""
        import ast

        tree = ast.parse("from os import path as p")
        visitor = ImportTrackingVisitor()
        visitor.visit(tree)
        assert "p" in visitor.imported_names
        assert "path" not in visitor.imported_names

    def test_skips_star_import(self) -> None:
        """Star imports are skipped."""
        import ast

        tree = ast.parse("from os import *")
        visitor = ImportTrackingVisitor()
        visitor.visit(tree)
        assert 0 == len(visitor.imported_names)

    def test_dotted_import_uses_first_component(self) -> None:
        """'import os.path' tracks 'os' as the available name."""
        import ast

        tree = ast.parse("import os.path")
        visitor = ImportTrackingVisitor()
        visitor.visit(tree)
        assert "os" in visitor.imported_names


class TestGetBaseClassName:
    """Tests for get_base_class_name."""

    def test_simple_name(self) -> None:
        """Simple Name base class returns the name."""
        arg = cst.Arg(value=cst.Name("Base"))
        assert "Base" == get_base_class_name(arg)

    def test_attribute_base(self) -> None:
        """Attribute base class returns the attr name."""
        arg = cst.Arg(
            value=cst.Attribute(value=cst.Name("mod"), attr=cst.Name("Base"))
        )
        assert "Base" == get_base_class_name(arg)

    def test_non_name_returns_none(self) -> None:
        """Non-Name/Attribute base returns None."""
        arg = cst.Arg(value=cst.Integer("42"))
        assert get_base_class_name(arg) is None


class TestHasDecorator:
    """Tests for has_decorator."""

    def test_simple_decorator(self) -> None:
        """Class with a simple decorator is detected."""
        code = "@dataclass\nclass Foo:\n    pass\n"
        module = cst.parse_module(code)
        class_node = module.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert has_decorator(class_node, "dataclass")

    def test_call_decorator(self) -> None:
        """Class with a call-style decorator is detected."""
        code = "@dataclass(frozen=True)\nclass Foo:\n    pass\n"
        module = cst.parse_module(code)
        class_node = module.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert has_decorator(class_node, "dataclass")

    def test_missing_decorator(self) -> None:
        """Class without the decorator returns False."""
        code = "class Foo:\n    pass\n"
        module = cst.parse_module(code)
        class_node = module.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert not has_decorator(class_node, "dataclass")

    def test_wrong_decorator(self) -> None:
        """Class with a different decorator returns False."""
        code = "@staticmethod\nclass Foo:\n    pass\n"
        module = cst.parse_module(code)
        class_node = module.body[0]
        assert isinstance(class_node, cst.ClassDef)
        assert not has_decorator(class_node, "dataclass")


class TestEvaluateExpressionAdvanced:
    """Additional tests for evaluate_expression."""

    def test_float(self) -> None:
        """Float node evaluates to truncated int."""
        assert 3 == evaluate_expression(cst.Float("3.9"))

    def test_binary_multiply(self) -> None:
        """Multiplication evaluates correctly."""
        node = cst.BinaryOperation(
            left=cst.Integer("3"),
            operator=cst.Multiply(),
            right=cst.Integer("4"),
        )
        assert 12 == evaluate_expression(node)

    def test_binary_subtract(self) -> None:
        """Subtraction evaluates correctly."""
        node = cst.BinaryOperation(
            left=cst.Integer("10"),
            operator=cst.Subtract(),
            right=cst.Integer("3"),
        )
        assert 7 == evaluate_expression(node)

    def test_binary_power(self) -> None:
        """Power evaluates correctly."""
        node = cst.BinaryOperation(
            left=cst.Integer("2"),
            operator=cst.Power(),
            right=cst.Integer("3"),
        )
        assert 8 == evaluate_expression(node)

    def test_unsupported_binary_op(self) -> None:
        """Unsupported binary operator returns None."""
        node = cst.BinaryOperation(
            left=cst.Integer("10"),
            operator=cst.FloorDivide(),
            right=cst.Integer("3"),
        )
        assert evaluate_expression(node) is None

    def test_nested_negative_unevaluable(self) -> None:
        """Unary minus on unevaluable returns None."""
        node = cst.UnaryOperation(
            operator=cst.Minus(),
            expression=cst.Name("x"),
        )
        assert evaluate_expression(node) is None


class TestEllipsisInCstNotTypes:
    """Tests for ellipsis_in_cst_not_types."""

    def test_no_ellipsis(self) -> None:
        """Module without ellipsis returns False."""
        module = cst.parse_module("x = 1\n")
        assert not ellipsis_in_cst_not_types(module)

    def test_ellipsis_in_assignment(self) -> None:
        """Ellipsis in assignment is invalid."""
        module = cst.parse_module("x = ...\n")
        assert ellipsis_in_cst_not_types(module)

    def test_ellipsis_in_type_annotation(self) -> None:
        """Ellipsis inside a type subscript is valid."""
        module = cst.parse_module("x: tuple[int, ...] = (1,)\n")
        assert not ellipsis_in_cst_not_types(module)

    def test_ellipsis_as_function_body(self) -> None:
        """Ellipsis as function body placeholder is valid."""
        module = cst.parse_module("def foo():\n    ...\n")
        assert not ellipsis_in_cst_not_types(module)


class TestAnyEllipsisInCst:
    """Tests for any_ellipsis_in_cst."""

    def test_no_ellipsis(self) -> None:
        """Module without ellipsis returns False."""
        module = cst.parse_module("x = 1\n")
        assert not any_ellipsis_in_cst(module)

    def test_has_ellipsis(self) -> None:
        """Module with any ellipsis returns True."""
        module = cst.parse_module("x = ...\n")
        assert any_ellipsis_in_cst(module)


class TestCompareUnparsedAstToSource:
    """Tests for compare_unparsed_ast_to_source."""

    def test_equivalent(self) -> None:
        """Equivalent code compares True."""
        assert compare_unparsed_ast_to_source("x = 1", "x  =  1")

    def test_different(self) -> None:
        """Different code compares False."""
        assert not compare_unparsed_ast_to_source("x = 1", "x = 2")


class TestUnparseParseSource:
    """Tests for unparse_parse_source."""

    def test_normalizes(self) -> None:
        """Normalizes formatting via ast round-trip."""
        assert "x = 1" == unparse_parse_source("x  =  1")


# ---------------------------------------------------------------------------
# _postprocess.py — additional coverage
# ---------------------------------------------------------------------------


class TestDocstringVisitor:
    """Tests for DocstringVisitor."""

    def test_collects_function_docstring(self) -> None:
        """Collects docstring from a top-level function."""
        module = cst.parse_module('def foo():\n    """Hello."""\n    pass\n')
        visitor = DocstringVisitor()
        module.visit(visitor)
        assert "Hello." == visitor.original_docstrings["foo"]

    def test_collects_class_docstring(self) -> None:
        """Collects docstring from a class."""
        module = cst.parse_module(
            'class Foo:\n    """Class doc."""\n    pass\n'
        )
        visitor = DocstringVisitor()
        module.visit(visitor)
        assert "Class doc." == visitor.original_docstrings["Foo"]

    def test_collects_method_docstring(self) -> None:
        """Collects docstring from a method with qualified name."""
        module = cst.parse_module(
            'class Foo:\n    def bar(self):\n        """Method doc."""\n        pass\n'
        )
        visitor = DocstringVisitor()
        module.visit(visitor)
        assert "Method doc." == visitor.original_docstrings["Foo.bar"]

    def test_no_docstring(self) -> None:
        """Function without docstring is not collected."""
        module = cst.parse_module("def foo():\n    pass\n")
        visitor = DocstringVisitor()
        module.visit(visitor)
        assert {} == visitor.original_docstrings

    def test_class_name_reset_after_leave(self) -> None:
        """class_name is reset to None after leaving a class."""
        module = cst.parse_module(
            "class Foo:\n    pass\ndef bar():\n    pass\n"
        )
        visitor = DocstringVisitor()
        module.visit(visitor)
        assert visitor.class_name is None


class TestDocstringTransformer:
    """Tests for DocstringTransformer."""

    def test_restores_missing_function_docstring(self) -> None:
        """Restores a removed function docstring."""
        docstrings = {"foo": "My docstring."}
        transformer = DocstringTransformer(docstrings)
        optimized = cst.parse_module("def foo():\n    pass\n")
        result = optimized.visit(transformer)
        assert "My docstring." in result.code

    def test_replaces_changed_function_docstring(self) -> None:
        """Replaces a modified function docstring with original."""
        docstrings = {"foo": "Original doc."}
        transformer = DocstringTransformer(docstrings)
        optimized = cst.parse_module(
            'def foo():\n    """Wrong doc."""\n    pass\n'
        )
        result = optimized.visit(transformer)
        assert "Original doc." in result.code

    def test_restores_missing_class_docstring(self) -> None:
        """Restores a removed class docstring."""
        docstrings = {"Foo": "Class docstring."}
        transformer = DocstringTransformer(docstrings)
        optimized = cst.parse_module("class Foo:\n    pass\n")
        result = optimized.visit(transformer)
        assert "Class docstring." in result.code

    def test_replaces_changed_class_docstring(self) -> None:
        """Replaces a modified class docstring with original."""
        docstrings = {"Foo": "Original class doc."}
        transformer = DocstringTransformer(docstrings)
        optimized = cst.parse_module(
            'class Foo:\n    """Wrong class doc."""\n    pass\n'
        )
        result = optimized.visit(transformer)
        assert "Original class doc." in result.code

    def test_no_change_without_original(self) -> None:
        """No change when no original docstring exists."""
        transformer = DocstringTransformer({})
        optimized = cst.parse_module("def foo():\n    pass\n")
        result = optimized.visit(transformer)
        assert result.code == optimized.code


class TestFixMissingDocstringAdvanced:
    """Additional tests for fix_missing_docstring."""

    def test_restores_class_docstring(self) -> None:
        """Restores a removed class docstring."""
        original = cst.parse_module(
            'class Foo:\n    """Class doc."""\n    pass\n'
        )
        optimized = cst.parse_module("class Foo:\n    pass\n")
        c = OptimizationCandidate(
            cst_module=optimized,
            explanation="faster",
            id="1",
        )
        result = fix_missing_docstring(original, [c])
        assert "Class doc." in result[0].cst_module.code

    def test_restores_method_docstring(self) -> None:
        """Restores a removed method docstring."""
        original = cst.parse_module(
            'class Foo:\n    def bar(self):\n        """Method doc."""\n        pass\n'
        )
        optimized = cst.parse_module(
            "class Foo:\n    def bar(self):\n        pass\n"
        )
        c = OptimizationCandidate(
            cst_module=optimized,
            explanation="faster",
            id="1",
        )
        result = fix_missing_docstring(original, [c])
        assert "Method doc." in result[0].cst_module.code


class TestDedupAndSortImports:
    """Tests for dedup_and_sort_imports."""

    def test_sorts_imports(self) -> None:
        """Unsorted imports get sorted."""
        code = "import sys\nimport os\n\nx = 1\n"
        module = cst.parse_module(code)
        c = OptimizationCandidate(
            cst_module=module,
            explanation="test",
            id="1",
        )
        result = dedup_and_sort_imports(module, [c])
        assert 1 == len(result)
        result_code = result[0].cst_module.code
        assert result_code.index("os") < result_code.index("sys")

    def test_no_change_when_sorted(self) -> None:
        """Already sorted imports are unchanged."""
        code = "import os\nimport sys\n"
        module = cst.parse_module(code)
        c = OptimizationCandidate(
            cst_module=module,
            explanation="test",
            id="1",
        )
        result = dedup_and_sort_imports(module, [c])
        assert 1 == len(result)


class TestEllipsisContainingCodeVisitor:
    """Tests for EllipsisContainingCodeVisitor."""

    def test_no_ellipsis(self) -> None:
        """Module without ellipsis sets flag to False."""
        module = cst.parse_module("x = 1\n")
        visitor = EllipsisContainingCodeVisitor()
        module.visit(visitor)
        assert not visitor.ellipsis_containing_code

    def test_has_ellipsis(self) -> None:
        """Module with ellipsis sets flag to True."""
        module = cst.parse_module("x = ...\n")
        visitor = EllipsisContainingCodeVisitor()
        module.visit(visitor)
        assert visitor.ellipsis_containing_code


class TestFilterEllipsisAdvanced:
    """Additional tests for filter_ellipsis_containing_code."""

    def test_keeps_non_ellipsis_candidate(self) -> None:
        """Candidate without ellipsis is kept when original has none."""
        module = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 2\n"),
            explanation="ok",
            id="1",
        )
        result = filter_ellipsis_containing_code(module, [c])
        assert 1 == len(result)

    def test_mixed_candidates(self) -> None:
        """Only candidates with ellipsis are filtered when original has none."""
        module = cst.parse_module("x = 1\n")
        c1 = OptimizationCandidate(
            cst_module=cst.parse_module("x = 2\n"),
            explanation="good",
            id="1",
        )
        c2 = OptimizationCandidate(
            cst_module=cst.parse_module("x = ...\n"),
            explanation="bad",
            id="2",
        )
        result = filter_ellipsis_containing_code(module, [c1, c2])
        assert 1 == len(result)
        assert "1" == result[0].id


class TestStripCommentsFromCode:
    """Tests for _strip_comments_from_code."""

    def test_removes_inline_comment(self) -> None:
        """Inline comments are removed."""
        code = "x = 1  # a comment\n"
        result = _strip_comments_from_code(code)
        assert "#" not in result
        assert "x = 1" in result

    def test_removes_full_line_comment(self) -> None:
        """Full-line comments are removed."""
        code = "# full line comment\nx = 1\n"
        result = _strip_comments_from_code(code)
        assert "full line comment" not in result
        assert "x = 1" in result

    def test_preserves_strings(self) -> None:
        """Hash inside strings is preserved."""
        code = 'x = "# not a comment"\n'
        result = _strip_comments_from_code(code)
        assert "# not a comment" in result

    def test_invalid_code_returns_original(self) -> None:
        """Invalid code returns the original string."""
        code = "def (:\n"
        assert code == _strip_comments_from_code(code)


class TestCleanExtraneousComments:
    """Tests for clean_extraneous_comments."""

    def test_preserves_original_comments(self) -> None:
        """Comments from original code are preserved in output."""
        original_code = "x = 1  # original comment\n"
        optimized_code = "x = 1  # original comment\n"
        original_module = cst.parse_module(original_code)
        optimized_module = cst.parse_module(optimized_code)
        result = clean_extraneous_comments(original_module, optimized_module)
        assert "original comment" in result.code

    def test_removes_new_comments_on_unchanged_lines(self) -> None:
        """New comments added to unchanged code lines are stripped."""
        original_code = "x = 1\ny = 2\n"
        optimized_code = "x = 1  # new comment\ny = 2\n"
        original_module = cst.parse_module(original_code)
        optimized_module = cst.parse_module(optimized_code)
        result = clean_extraneous_comments(original_module, optimized_module)
        assert "new comment" not in result.code


class TestCleanExtraneousCommentsPipeline:
    """Tests for clean_extraneous_comments_pipeline."""

    def test_processes_candidates(self) -> None:
        """Pipeline processes all candidates."""
        original = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 2  # optimized\n"),
            explanation="test",
            id="1",
        )
        result = clean_extraneous_comments_pipeline(original, [c])
        assert 1 == len(result)


class TestEqualityCheckAdvanced:
    """Additional edge-case tests for equality_check."""

    def test_with_explicit_original_code(self) -> None:
        """Original code can be passed explicitly."""
        module = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 2\n"),
            explanation="different",
            id="1",
        )
        result = equality_check(module, [c], original_code="x = 1\n")
        assert 1 == len(result)

    def test_filters_with_explicit_code(self) -> None:
        """Identical candidate is filtered with explicit code."""
        module = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 1\n"),
            explanation="same",
            id="1",
        )
        result = equality_check(module, [c], original_code="x = 1\n")
        assert 0 == len(result)

    def test_unparseable_original_falls_back_to_string(self) -> None:
        """When original is unparseable, falls back to string comparison."""
        module = cst.parse_module("x = 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 2\n"),
            explanation="different",
            id="1",
        )
        result = equality_check(module, [c], original_code="def (:\n")
        assert 1 == len(result)


class TestExtractNamesFromCstAnnotation:
    """Tests for extract_names_from_cst_annotation."""

    def test_simple_name(self) -> None:
        """Extracts simple name from annotation."""
        names: set[str] = set()
        extract_names_from_cst_annotation(cst.Name("int"), names)
        assert {"int"} == names

    def test_attribute(self) -> None:
        """Extracts base name from attribute annotation."""
        names: set[str] = set()
        node = cst.Attribute(value=cst.Name("typing"), attr=cst.Name("List"))
        extract_names_from_cst_annotation(node, names)
        assert "typing" in names

    def test_subscript(self) -> None:
        """Extracts names from subscript annotation like List[int]."""
        names: set[str] = set()
        node = cst.Subscript(
            value=cst.Name("List"),
            slice=[
                cst.SubscriptElement(slice=cst.Index(value=cst.Name("int")))
            ],
        )
        extract_names_from_cst_annotation(node, names)
        assert "List" in names
        assert "int" in names

    def test_binary_operation(self) -> None:
        """Extracts names from union-style binary annotation."""
        names: set[str] = set()
        node = cst.BinaryOperation(
            left=cst.Name("str"),
            operator=cst.BitOr(),
            right=cst.Name("int"),
        )
        extract_names_from_cst_annotation(node, names)
        assert {"str", "int"} == names

    def test_tuple_annotation(self) -> None:
        """Extracts names from tuple annotation."""
        names: set[str] = set()
        node = cst.Tuple(
            elements=[
                cst.Element(value=cst.Name("int")),
                cst.Element(value=cst.Name("str")),
            ]
        )
        extract_names_from_cst_annotation(node, names)
        assert {"int", "str"} == names


class TestCSTAnnotationNameCollector:
    """Tests for CSTAnnotationNameCollector."""

    def test_collects_function_annotations(self) -> None:
        """Collects annotation names from function parameters and return."""
        module = cst.parse_module(
            "def foo(x: MyType) -> ReturnType:\n    pass\n"
        )
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        assert "MyType" in collector.annotation_names
        assert "ReturnType" in collector.annotation_names

    def test_collects_imports(self) -> None:
        """Tracks imported names."""
        module = cst.parse_module("import os\nfrom sys import path\n")
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        assert "os" in collector.imported_names
        assert "path" in collector.imported_names

    def test_collects_class_names(self) -> None:
        """Tracks defined class names."""
        module = cst.parse_module("class Foo:\n    pass\n")
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        assert "Foo" in collector.defined_names

    def test_collects_function_names(self) -> None:
        """Tracks defined function names."""
        module = cst.parse_module("def foo():\n    pass\n")
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        assert "foo" in collector.defined_names

    def test_undefined_annotations(self) -> None:
        """Returns annotation names that are not defined or imported."""
        module = cst.parse_module("def foo(x: MyType) -> None:\n    pass\n")
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        undefined = collector.get_undefined_annotation_names()
        assert "MyType" in undefined
        assert "None" not in undefined

    def test_no_undefined_with_import(self) -> None:
        """No undefined names when all annotations are imported."""
        module = cst.parse_module(
            "from mymod import MyType\ndef foo(x: MyType) -> None:\n    pass\n"
        )
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        undefined = collector.get_undefined_annotation_names()
        assert "MyType" not in undefined

    def test_collects_kwonly_params(self) -> None:
        """Collects annotations from keyword-only parameters."""
        module = cst.parse_module("def foo(*, x: KWType) -> None:\n    pass\n")
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        assert "KWType" in collector.annotation_names

    def test_collects_star_arg_annotation(self) -> None:
        """Collects annotations from *args parameter."""
        module = cst.parse_module(
            "def foo(*args: ArgsType) -> None:\n    pass\n"
        )
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        assert "ArgsType" in collector.annotation_names

    def test_collects_star_kwarg_annotation(self) -> None:
        """Collects annotations from **kwargs parameter."""
        module = cst.parse_module(
            "def foo(**kwargs: KwargsType) -> None:\n    pass\n"
        )
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        assert "KwargsType" in collector.annotation_names

    def test_collects_ann_assign(self) -> None:
        """Collects annotation names from annotated assignments."""
        module = cst.parse_module("x: MyType = 1\n")
        collector = CSTAnnotationNameCollector()
        module.visit(collector)
        assert "MyType" in collector.annotation_names


class TestFixForwardReferences:
    """Tests for fix_forward_references."""

    def test_adds_import_when_needed(self) -> None:
        """Adds future annotations import for undefined annotation names."""
        module = cst.parse_module("x = 1\n")
        code = "def foo(x: MyType) -> None:\n    pass\n"
        c = OptimizationCandidate(
            cst_module=cst.parse_module(code),
            explanation="test",
            id="1",
        )
        result = fix_forward_references(module, [c])
        assert 1 == len(result)
        assert (
            "from __future__ import annotations" in result[0].cst_module.code
        )

    def test_no_change_when_all_defined(self) -> None:
        """No change when all annotations are builtins."""
        module = cst.parse_module("x = 1\n")
        code = "def foo(x: int) -> None:\n    pass\n"
        c = OptimizationCandidate(
            cst_module=cst.parse_module(code),
            explanation="test",
            id="1",
        )
        result = fix_forward_references(module, [c])
        assert 1 == len(result)
        assert (
            "from __future__ import annotations"
            not in result[0].cst_module.code
        )

    def test_no_change_when_already_present(self) -> None:
        """No change when future annotations import already exists."""
        module = cst.parse_module("x = 1\n")
        code = "from __future__ import annotations\ndef foo(x: MyType) -> None:\n    pass\n"
        c = OptimizationCandidate(
            cst_module=cst.parse_module(code),
            explanation="test",
            id="1",
        )
        result = fix_forward_references(module, [c])
        assert 1 == len(result)


class TestPostprocessingPipelineAdvanced:
    """Additional tests for optimizations_postprocessing_pipeline."""

    def test_pipeline_restores_docstring(self) -> None:
        """Pipeline restores removed docstrings."""
        original = cst.parse_module(
            'def foo():\n    """My doc."""\n    return 1\n'
        )
        c = OptimizationCandidate(
            cst_module=cst.parse_module("def foo():\n    return 2\n"),
            explanation="faster",
            id="1",
        )
        result = optimizations_postprocessing_pipeline(original, [c])
        assert 1 == len(result)
        assert "My doc." in result[0].cst_module.code

    def test_pipeline_adds_future_annotations(self) -> None:
        """Pipeline adds future annotations when needed."""
        original = cst.parse_module("def foo():\n    return 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module(
                "def foo(x: MyCustomType) -> None:\n    return 2\n"
            ),
            explanation="typed",
            id="1",
        )
        result = optimizations_postprocessing_pipeline(original, [c])
        assert 1 == len(result)
        assert (
            "from __future__ import annotations" in result[0].cst_module.code
        )

    def test_pipeline_cleans_explanations(self) -> None:
        """Pipeline cleans up LLM explanation artifacts."""
        original = cst.parse_module("x = 1 + 1\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = 2\n"),
            explanation="```python\nx = 2\n```",
            id="1",
        )
        result = optimizations_postprocessing_pipeline(original, [c])
        assert 1 == len(result)
        assert "```" not in result[0].explanation

    def test_pipeline_filters_ellipsis(self) -> None:
        """Pipeline filters candidates that introduce ellipsis."""
        original = cst.parse_module("x = 1\ny = 2\n")
        c = OptimizationCandidate(
            cst_module=cst.parse_module("x = ...\ny = ...\n"),
            explanation="bad",
            id="1",
        )
        result = optimizations_postprocessing_pipeline(original, [c])
        assert 0 == len(result)
