"""Dependency collection and unused definition removal via libcst."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import libcst as cst

from .models import UsageInfo

if TYPE_CHECKING:
    from collections.abc import Callable

log = logging.getLogger(__name__)


def extract_names_from_targets(target: cst.CSTNode) -> list[str]:
    """Extract variable names from a target, including tuple unpacking."""
    names = []

    # Handle a simple name
    if isinstance(target, cst.Name):
        names.append(target.value)

    # Handle any node with a value attribute (StarredElement, etc.)
    elif hasattr(target, "value"):
        names.extend(extract_names_from_targets(target.value))

    # Handle any node with elements attribute (tuples, lists, etc.)
    elif hasattr(target, "elements"):
        for element in target.elements:
            # Recursive call for each element
            names.extend(extract_names_from_targets(element))

    return names


def is_assignment_used(
    node: cst.CSTNode, definitions: dict[str, UsageInfo], name_prefix: str = ""
) -> bool:
    """Return True if any name from the assignment is used by the target function."""
    if isinstance(node, cst.Assign):
        targets = [target.target for target in node.targets]
    elif isinstance(node, (cst.AnnAssign, cst.AugAssign)):
        targets = [node.target]
    else:
        return False
    for target in targets:
        for name in extract_names_from_targets(target):
            lookup = f"{name_prefix}{name}" if name_prefix else name
            if (
                lookup in definitions
                and definitions[lookup].used_by_qualified_function
            ):
                return True
    return False


def recurse_sections(  # noqa: C901, PLR0912
    node: cst.CSTNode,
    section_names: list[str],
    prune_fn: Callable[[cst.CSTNode], tuple[cst.CSTNode | None, bool]],
    *,
    keep_non_target_children: bool = False,
) -> tuple[cst.CSTNode | None, bool]:
    """Recursively prune sections of a CST node, keeping only target subtrees."""
    updates: dict[str, list[cst.CSTNode] | cst.CSTNode] = {}
    found_any_target = False
    for section in section_names:
        original_content = getattr(node, section, None)
        if isinstance(original_content, (list, tuple)):
            new_children = []
            section_found_target = False
            for child in original_content:
                filtered, found_target = prune_fn(child)
                if filtered:
                    new_children.append(filtered)
                section_found_target |= found_target
            if keep_non_target_children:
                if section_found_target or new_children:
                    found_any_target |= section_found_target
                    updates[section] = new_children
            elif section_found_target:
                found_any_target = True
                updates[section] = new_children
        elif original_content is not None:
            filtered, found_target = prune_fn(original_content)
            if keep_non_target_children:
                found_any_target |= found_target
                if filtered:
                    updates[section] = filtered
            elif found_target:
                found_any_target = True
                if filtered:
                    updates[section] = filtered
    if keep_non_target_children:
        if updates:
            return node.with_changes(**updates), found_any_target
        return None, False
    if not found_any_target:
        return None, False
    return (node.with_changes(**updates) if updates else node), True


def collect_top_level_definitions(  # noqa: C901, PLR0912
    node: cst.CSTNode,
    definitions: dict[str, UsageInfo] | None = None,
) -> dict[str, UsageInfo]:
    """Collect all top-level definitions from a CST node."""
    if definitions is None:
        definitions = {}

    if isinstance(node, cst.FunctionDef):
        name = node.name.value
        definitions[name] = UsageInfo(name=name)
        return definitions

    if isinstance(node, cst.ClassDef):
        name = node.name.value
        definitions[name] = UsageInfo(name=name)
        if isinstance(node.body, cst.IndentedBlock):
            prefix = name + "."
            for statement in node.body.body:
                if isinstance(statement, cst.FunctionDef):
                    method_name = prefix + statement.name.value
                    definitions[method_name] = UsageInfo(name=method_name)
        return definitions

    if isinstance(node, cst.Assign):
        for target in node.targets:
            for name in extract_names_from_targets(target.target):
                definitions[name] = UsageInfo(name=name)
        return definitions

    if isinstance(node, (cst.AnnAssign, cst.AugAssign)):
        for name in extract_names_from_targets(node.target):
            definitions[name] = UsageInfo(name=name)
        return definitions

    for section in get_section_names(node):
        original_content = getattr(node, section, None)
        if isinstance(original_content, (list, tuple)):
            for child in original_content:
                collect_top_level_definitions(child, definitions)
        elif original_content is not None:
            collect_top_level_definitions(original_content, definitions)

    return definitions


def get_section_names(node: cst.CSTNode) -> list[str]:
    """Return section attribute names for a given node."""
    possible_sections = ["body", "orelse", "finalbody", "handlers"]
    return [sec for sec in possible_sections if hasattr(node, sec)]


class DependencyCollector(cst.CSTVisitor):
    """Collect dependencies between definitions via depth tracking."""

    def __init__(self, definitions: dict[str, UsageInfo]) -> None:
        """Initialize with the definitions dictionary to populate with dependencies."""
        super().__init__()
        self.definitions = definitions
        # Track function and class depths
        self.function_depth = 0
        self.class_depth = 0
        # Track top-level qualified names
        self.current_top_level_name = ""
        self.current_class = ""
        # Track if we're processing a top-level variable
        self.processing_variable = False
        self.current_variable_names: set[str] = set()
        # Track Name nodes that are the .attr part of Attribute nodes (by id)
        self.attr_name_ids: set[int] = set()

    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:  # noqa: N802
        """Track entering a function definition and set current context."""
        function_name = node.name.value

        if self.function_depth == 0:
            # This is a top-level function
            if self.class_depth > 0:
                # Inside a class: track deps at class level
                self.current_top_level_name = (
                    f"{self.current_class}.{function_name}"
                )
            else:
                # Regular top-level function
                self.current_top_level_name = function_name

        for param in node.params.params:
            if param.annotation:
                self._extract_names_from_annotation(
                    param.annotation.annotation
                )

        self.function_depth += 1

    def _extract_names_from_annotation(self, node: cst.CSTNode) -> None:
        """Record type annotation names as dependencies of the current definition."""
        if isinstance(node, cst.Name):
            name = node.value
            if (
                name in self.definitions
                and name != self.current_top_level_name
                and self.current_top_level_name
            ):
                self.definitions[self.current_top_level_name].dependencies.add(
                    name
                )
        elif isinstance(node, cst.Subscript):
            self._extract_names_from_annotation(node.value)
            for slice_item in node.slice:
                if hasattr(slice_item, "slice"):
                    self._extract_names_from_annotation(slice_item.slice)
        elif isinstance(node, cst.Attribute):
            self._extract_names_from_annotation(node.value)

    def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:  # noqa: N802
        """Track exiting a function definition and restore context."""
        self.function_depth -= 1

        if self.function_depth == 0 and self.class_depth == 0:
            # Exiting top-level function that's not in a class
            self.current_top_level_name = ""

    def visit_ClassDef(self, node: cst.ClassDef) -> None:  # noqa: N802
        """Track entering a class definition and record base class dependencies."""
        class_name = node.name.value

        if self.class_depth == 0:
            # This is a top-level class
            self.current_class = class_name
            self.current_top_level_name = class_name

            # Track base classes as dependencies
            for base in node.bases:
                if isinstance(base.value, cst.Name):
                    base_name = base.value.value
                    if (
                        base_name in self.definitions
                        and class_name in self.definitions
                    ):
                        self.definitions[class_name].dependencies.add(
                            base_name
                        )
                elif isinstance(base.value, cst.Attribute):
                    # Handle cases like module.ClassName
                    attr_name = base.value.attr.value
                    if (
                        attr_name in self.definitions
                        and class_name in self.definitions
                    ):
                        self.definitions[class_name].dependencies.add(
                            attr_name
                        )

        self.class_depth += 1

    def leave_ClassDef(self, original_node: cst.ClassDef) -> None:  # noqa: N802
        """Track exiting a class definition and restore context."""
        self.class_depth -= 1

        if self.class_depth == 0:
            # Exiting top-level class
            self.current_class = ""
            self.current_top_level_name = ""

    def visit_Assign(self, node: cst.Assign) -> None:  # noqa: N802
        """Track top-level assignment targets as current context for dependency collection."""
        # Only handle top-level assignments
        if self.function_depth == 0 and self.class_depth == 0:
            for target in node.targets:
                # Extract all variable names from the target
                names = extract_names_from_targets(target.target)

                # Check if any names are tracked top-level defs
                tracked_names = [
                    name for name in names if name in self.definitions
                ]
                if tracked_names:
                    self.processing_variable = True
                    self.current_variable_names.update(tracked_names)
                    # Use first tracked name as current top-level
                    self.current_top_level_name = tracked_names[0]

    def leave_Assign(self, original_node: cst.Assign) -> None:  # noqa: N802
        """Reset assignment processing state after leaving the node."""
        if self.processing_variable:
            self.processing_variable = False
            self.current_variable_names.clear()
            self.current_top_level_name = ""

    def visit_AnnAssign(self, node: cst.AnnAssign) -> None:  # noqa: N802
        """Record annotation dependencies for annotated assignments."""
        self.processing_variable = True
        if isinstance(node.target, cst.Name):
            self.current_variable_names.add(node.target.value)
        else:
            self.current_variable_names.update(
                extract_names_from_targets(node.target)
            )

        self._extract_names_from_annotation(node.annotation.annotation)

        self.processing_variable = False
        self.current_variable_names.clear()

    def visit_Attribute(self, node: cst.Attribute) -> None:  # noqa: N802
        """Mark attribute names to exclude from dependency tracking."""
        self.attr_name_ids.add(id(node.attr))

    def leave_Attribute(self, original_node: cst.Attribute) -> None:  # noqa: N802
        """Unmark attribute name after leaving the node."""
        self.attr_name_ids.discard(id(original_node.attr))

    def visit_Name(self, node: cst.Name) -> None:  # noqa: N802
        """Record a name reference as a dependency of the current definition."""
        name = node.value

        # Skip if we're not inside a tracked definition
        if (
            not self.current_top_level_name
            or self.current_top_level_name not in self.definitions
        ):
            return

        # Skip if we're looking at the variable name itself in an assignment
        if self.processing_variable and name in self.current_variable_names:
            return

        if name in self.definitions and name != self.current_top_level_name:
            # Skip .attr part of Attribute (e.g., 'x' in 'self.x')
            # Only track base/value, not attribute name
            if self.class_depth > 0:
                if id(node) in self.attr_name_ids:
                    return
                # .value (base): only skip self/cls
                if name in ("self", "cls"):
                    return
            self.definitions[self.current_top_level_name].dependencies.add(
                name
            )


class QualifiedFunctionUsageMarker:
    """Marks definitions that are used by specific qualified functions."""

    def __init__(
        self,
        definitions: dict[str, UsageInfo],
        qualified_function_names: set[str],
    ) -> None:
        """Initialize with definitions and the set of target qualified function names."""
        self.definitions = definitions
        self.qualified_function_names = qualified_function_names
        self.expanded_qualified_functions = self._expand_qualified_functions()

    def _expand_qualified_functions(self) -> set[str]:
        """Expand the qualified function names to include related methods."""
        expanded = set(self.qualified_function_names)

        # Add containing classes and dunder methods
        for qualified_name in list(self.qualified_function_names):
            if "." in qualified_name:
                class_name, _method_name = qualified_name.split(".", 1)

                # Add the class itself
                expanded.add(class_name)

                # Add all dunder methods of the class
                for name in self.definitions:
                    if name.startswith(f"{class_name}.__") and name.endswith(
                        "__"
                    ):
                        expanded.add(name)

        return expanded

    def mark_used_definitions(self) -> None:
        """Mark qualified functions and their deps as used."""
        defs = self.definitions
        for func_name in self.expanded_qualified_functions & defs.keys():
            defs[func_name].used_by_qualified_function = True
            for dep in defs[func_name].dependencies:
                self.mark_as_used_recursively(dep)

    def mark_as_used_recursively(self, name: str) -> None:
        """Mark a name and all its dependencies as used recursively."""
        if name not in self.definitions:
            return

        if self.definitions[name].used_by_qualified_function:
            return  # Already marked

        self.definitions[name].used_by_qualified_function = True

        # Mark all dependencies as used
        for dep in self.definitions[name].dependencies:
            self.mark_as_used_recursively(dep)


def remove_unused_definitions_recursively(  # noqa: C901, PLR0911
    node: cst.CSTNode, definitions: dict[str, UsageInfo]
) -> tuple[cst.CSTNode | None, bool]:
    """Recursively filter the node to remove unused definitions.

    Returns (filtered_node_or_None, used_by_function).
    """
    # Skip import statements
    if isinstance(node, (cst.Import, cst.ImportFrom)):
        return node, True

    # Never remove function definitions
    if isinstance(node, cst.FunctionDef):
        return node, True

    if isinstance(node, cst.ClassDef):
        class_name = node.name.value
        class_has_dependencies = (
            class_name in definitions
            and definitions[class_name].used_by_qualified_function
        )

        if isinstance(node.body, cst.IndentedBlock):
            new_statements: list[cst.BaseStatement] = []
            for statement in node.body.body:
                if isinstance(statement, cst.FunctionDef):
                    new_statements.append(statement)
                elif isinstance(
                    statement, (cst.Assign, cst.AnnAssign, cst.AugAssign)
                ):
                    if class_has_dependencies or is_assignment_used(
                        statement, definitions, name_prefix=f"{class_name}."
                    ):
                        new_statements.append(statement)
                else:
                    new_statements.append(statement)
            return node.with_changes(
                body=node.body.with_changes(body=new_statements)
            ), True

        return node, class_has_dependencies

    # Handle assignments (Assign, AnnAssign, AugAssign)
    if isinstance(node, (cst.Assign, cst.AnnAssign, cst.AugAssign)):
        if is_assignment_used(node, definitions):
            return node, True
        return None, False

    # For other nodes, recursively process children
    section_names = get_section_names(node)
    if not section_names:
        return node, False
    return recurse_sections(
        node,
        section_names,
        lambda child: remove_unused_definitions_recursively(
            child, definitions
        ),
    )


def collect_top_level_defs_with_dependencies(
    code: str | cst.Module,
) -> dict[str, UsageInfo]:
    """Collect definitions and their dependencies (CST pass).

    Returns defs with dependencies but no usage marks set.
    Reuse across multiple mark_defs_for_functions calls
    to skip the DependencyCollector traversal.
    """
    module = code if isinstance(code, cst.Module) else cst.parse_module(code)
    definitions = collect_top_level_definitions(module)
    dependency_collector = DependencyCollector(definitions)
    module.visit(dependency_collector)
    return definitions


def mark_defs_for_functions(
    base_defs: dict[str, UsageInfo], qualified_function_names: set[str]
) -> dict[str, UsageInfo]:
    """Copy *base_defs* with usage marks for *qualified_function_names*.

    Cheap (dict copy + graph walk), reusable with different
    function name sets without re-traversing the CST.
    """
    marked = {
        k: UsageInfo(name=v.name, dependencies=v.dependencies)
        for k, v in base_defs.items()
    }
    usage_marker = QualifiedFunctionUsageMarker(
        marked, qualified_function_names
    )
    usage_marker.mark_used_definitions()
    return marked


def collect_top_level_defs_with_usages(
    code: str | cst.Module,
    qualified_function_names: set[str],
) -> dict[str, UsageInfo]:
    """Collect definitions and mark usages in one pass."""
    base_defs = collect_top_level_defs_with_dependencies(code)
    return mark_defs_for_functions(base_defs, qualified_function_names)


def remove_unused_definitions_by_function_names(
    code: str | cst.Module,
    qualified_function_names: set[str],
    defs_with_usages: dict[str, UsageInfo] | None = None,
) -> cst.Module:
    """Remove top-level definitions not used by *qualified_function_names*."""
    try:
        module = (
            code if isinstance(code, cst.Module) else cst.parse_module(code)
        )
    except Exception:  # noqa: BLE001
        log.debug(
            "Failed to parse code with libcst",
            exc_info=True,
        )
        return code if isinstance(code, cst.Module) else cst.parse_module("")

    try:
        if defs_with_usages is None:
            base = collect_top_level_defs_with_dependencies(
                module,
            )
            defs_with_usages = mark_defs_for_functions(
                base,
                qualified_function_names,
            )

        result, _ = remove_unused_definitions_recursively(
            module,
            defs_with_usages,
        )
        if not isinstance(result, cst.Module):
            return cst.parse_module("")
        return result  # noqa: TRY300
    except Exception:  # noqa: BLE001
        log.debug(
            "Error removing unused definitions",
            exc_info=True,
        )
        return module
