"""Tests for _context.dependencies (CST dependency collection)."""

from __future__ import annotations

import libcst as cst

from codeflash_python.context.dependencies import (
    UsageInfo,
    collect_top_level_defs_with_dependencies,
    extract_names_from_targets,
    get_section_names,
    is_assignment_used,
    mark_defs_for_functions,
    remove_unused_definitions_by_function_names,
    remove_unused_definitions_recursively,
)


class TestExtractNamesFromTargets:
    """Tests for extract_names_from_targets."""

    def test_simple_name(self) -> None:
        """A single Name node yields its value."""
        node = cst.parse_expression("x")
        assert ["x"] == extract_names_from_targets(node)

    def test_tuple_unpacking(self) -> None:
        """Tuple targets yield all contained names."""
        node = cst.parse_expression("(a, b, c)")
        assert ["a", "b", "c"] == extract_names_from_targets(
            node,
        )

    def test_starred_element(self) -> None:
        """StarredElement targets yield the inner name."""
        stmt = cst.parse_statement("*rest, = items\n")
        assert isinstance(stmt, cst.SimpleStatementLine)
        assign = stmt.body[0]
        assert isinstance(assign, cst.Assign)
        target = assign.targets[0].target
        result = extract_names_from_targets(target)
        assert "rest" in result


class TestIsAssignmentUsed:
    """Tests for is_assignment_used."""

    def test_used_assignment(self) -> None:
        """An Assign whose target is marked used returns True."""
        stmt = cst.parse_statement("x = 1\n")
        assert isinstance(stmt, cst.SimpleStatementLine)
        node = stmt.body[0]
        defs = {
            "x": UsageInfo(
                name="x",
                used_by_qualified_function=True,
            ),
        }
        assert is_assignment_used(node, defs) is True

    def test_unused_assignment(self) -> None:
        """An Assign whose target is not marked returns False."""
        stmt = cst.parse_statement("y = 2\n")
        assert isinstance(stmt, cst.SimpleStatementLine)
        node = stmt.body[0]
        defs = {
            "y": UsageInfo(name="y"),
        }
        assert is_assignment_used(node, defs) is False

    def test_with_prefix(self) -> None:
        """Name prefix is prepended for class-level lookups."""
        stmt = cst.parse_statement("val = 3\n")
        assert isinstance(stmt, cst.SimpleStatementLine)
        node = stmt.body[0]
        defs = {
            "Cls.val": UsageInfo(
                name="Cls.val",
                used_by_qualified_function=True,
            ),
        }
        assert (
            is_assignment_used(
                node,
                defs,
                name_prefix="Cls.",
            )
            is True
        )


class TestGetSectionNames:
    """Tests for get_section_names."""

    def test_module_has_body(self) -> None:
        """A Module node has a 'body' section."""
        module = cst.parse_module("")
        assert "body" in get_section_names(module)

    def test_name_has_no_sections(self) -> None:
        """A Name node has no section attributes."""
        node = cst.parse_expression("x")
        assert [] == get_section_names(node)


class TestCollectTopLevelDefsWithDependencies:
    """Tests for collect_top_level_defs_with_dependencies."""

    def test_function_collected(self) -> None:
        """A top-level function is collected."""
        code = "def greet(): pass\n"
        defs = collect_top_level_defs_with_dependencies(code)
        assert "greet" in defs

    def test_class_and_methods(self) -> None:
        """A class and its methods are collected."""
        code = "class Foo:\n    def bar(self): pass\n"
        defs = collect_top_level_defs_with_dependencies(code)
        assert "Foo" in defs
        assert "Foo.bar" in defs

    def test_variable_collected(self) -> None:
        """A top-level assignment is collected."""
        code = "X = 42\n"
        defs = collect_top_level_defs_with_dependencies(code)
        assert "X" in defs

    def test_dependency_tracked(self) -> None:
        """A function referencing another records a dependency."""
        code = "def helper(): pass\ndef caller(): return helper()\n"
        defs = collect_top_level_defs_with_dependencies(code)
        assert "helper" in defs["caller"].dependencies

    def test_class_base_dependency(self) -> None:
        """A class base class is recorded as a dependency."""
        code = "class Base: pass\nclass Child(Base): pass\n"
        defs = collect_top_level_defs_with_dependencies(code)
        assert "Base" in defs["Child"].dependencies


class TestMarkDefsForFunctions:
    """Tests for mark_defs_for_functions."""

    def test_marks_target_and_deps(self) -> None:
        """The target and its transitive deps are marked."""
        code = (
            "def a(): return b()\n"
            "def b(): return c()\n"
            "def c(): pass\n"
            "def unrelated(): pass\n"
        )
        base = collect_top_level_defs_with_dependencies(code)
        marked = mark_defs_for_functions(base, {"a"})
        assert marked["a"].used_by_qualified_function is True
        assert marked["b"].used_by_qualified_function is True
        assert marked["c"].used_by_qualified_function is True
        assert marked["unrelated"].used_by_qualified_function is False

    def test_base_defs_unchanged(self) -> None:
        """Marking does not mutate the original base_defs."""
        code = "def f(): pass\n"
        base = collect_top_level_defs_with_dependencies(code)
        mark_defs_for_functions(base, {"f"})
        assert base["f"].used_by_qualified_function is False

    def test_class_method_marks_class(self) -> None:
        """Marking a class method also marks the class."""
        code = "class Cls:\n    def method(self): pass\n"
        base = collect_top_level_defs_with_dependencies(code)
        marked = mark_defs_for_functions(
            base,
            {"Cls.method"},
        )
        assert marked["Cls"].used_by_qualified_function is True
        assert marked["Cls.method"].used_by_qualified_function is True


class TestRemoveUnusedDefinitions:
    """Tests for remove_unused_definitions_recursively."""

    def test_keeps_imports(self) -> None:
        """Import statements are always kept."""
        module = cst.parse_module("import os\n")
        node = module.body[0]
        result, used = remove_unused_definitions_recursively(
            node,
            {},
        )
        assert result is not None
        assert used is True

    def test_keeps_function_defs(self) -> None:
        """Function definitions are always kept."""
        module = cst.parse_module("def f(): pass\n")
        node = module.body[0]
        result, used = remove_unused_definitions_recursively(
            node,
            {},
        )
        assert result is not None
        assert used is True

    def test_removes_unused_assignment(self) -> None:
        """Unused assignments are removed."""
        stmt = cst.parse_statement("x = 1\n")
        assert isinstance(stmt, cst.SimpleStatementLine)
        node = stmt.body[0]
        defs = {"x": UsageInfo(name="x")}
        result, used = remove_unused_definitions_recursively(
            node,
            defs,
        )
        assert result is None
        assert used is False


class TestRemoveUnusedByFunctionNames:
    """Tests for remove_unused_definitions_by_function_names."""

    def test_keeps_used_removes_unused(self) -> None:
        """Used definitions are kept, unused are removed."""
        code = (
            "import os\n"
            "USED = 1\n"
            "UNUSED = 2\n"
            "def target(): return USED\n"
            "def other(): pass\n"
        )
        result = remove_unused_definitions_by_function_names(
            code,
            {"target"},
        )
        output = result.code
        assert "USED" in output
        assert "UNUSED" not in output
        assert "def target" in output
        assert "def other" in output
        assert "import os" in output

    def test_invalid_code_returns_empty(self) -> None:
        """Invalid code returns an empty module."""
        result = remove_unused_definitions_by_function_names(
            "def (\n",
            {"f"},
        )
        assert "" == result.code.strip()

    def test_with_precomputed_defs(self) -> None:
        """Pre-computed defs_with_usages can be passed in."""
        code = "X = 1\ndef f(): return X\n"
        base = collect_top_level_defs_with_dependencies(code)
        marked = mark_defs_for_functions(base, {"f"})
        result = remove_unused_definitions_by_function_names(
            code,
            {"f"},
            defs_with_usages=marked,
        )
        assert "X = 1" in result.code
