from __future__ import annotations

from dataclasses import dataclass, field

import libcst as cst

from core.languages.python.cst_utils import DepthTrackingMixin
from core.languages.python.testgen.ast_utils.test_detection import is_test_class, is_test_function_name


@dataclass
class UsageInfo:
    """Information about a name and its usage."""

    name: str
    used_by_test: bool = False
    dependencies: set[str] = field(default_factory=set)


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

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

    # Handle StarredElement (has a value attribute)
    elif isinstance(target, cst.StarredElement):
        names.extend(extract_names_from_targets(target.value))

    # Handle tuples and lists (have elements attribute)
    elif isinstance(target, cst.Tuple | cst.List):
        for element in target.elements:
            # Element could be Element, StarredElement, etc.
            if isinstance(element, (cst.Element, cst.StarredElement)):
                names.extend(extract_names_from_targets(element.value))

    return names


def collect_top_level_definitions(
    node: cst.CSTNode, definitions: dict[str, UsageInfo] | None = None
) -> dict[str, UsageInfo]:
    """Recursively collect all top-level variable, function, and class definitions."""
    if definitions is None:
        definitions = {}

    # Handle top-level function definitions
    if isinstance(node, cst.FunctionDef):
        name = node.name.value
        definitions[name] = UsageInfo(
            name=name,
            used_by_test=is_test_function_name(name),  # Test functions are inherently "used"
        )
        return definitions

    # Handle top-level class definitions
    if isinstance(node, cst.ClassDef):
        name = node.name.value
        definitions[name] = UsageInfo(name=name, used_by_test=is_test_class(node))
        return definitions

    # Handle top-level variable assignments
    if isinstance(node, cst.Assign):
        for target in node.targets:
            names = extract_names_from_targets(target.target)
            for name in names:
                definitions[name] = UsageInfo(name=name)
        return definitions

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

    # Recursively process children. Takes care of top level assignments in if/else/while/for blocks
    section_names = get_section_names(node)

    if section_names:
        for section in section_names:
            original_content = getattr(node, section, None)
            # If section contains a list of nodes
            if isinstance(original_content, list | tuple):
                for child in original_content:
                    collect_top_level_definitions(child, definitions)
            # If section contains a single node
            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 the section attribute names (e.g., body, orelse) for a given node if they exist."""
    possible_sections = ["body", "orelse", "finalbody", "handlers"]
    return [sec for sec in possible_sections if hasattr(node, sec)]


class DependencyCollector(DepthTrackingMixin, cst.CSTVisitor):
    """Collects dependencies between definitions using the visitor pattern with depth tracking."""

    def __init__(self, definitions: dict[str, UsageInfo]) -> None:
        DepthTrackingMixin.__init__(self)
        cst.CSTVisitor.__init__(self)
        self.definitions = definitions
        # 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()

    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
        function_name = node.name.value

        if self._function_depth == 0:
            # This is a top-level function
            if self._is_inside_class():
                # If inside a class, we're now tracking dependencies at the class level
                self.current_top_level_name = self.current_class
            else:
                # Regular top-level function
                self.current_top_level_name = function_name

        self._visit_function()

    def leave_FunctionDef(self, original_node: cst.FunctionDef) -> None:  # noqa: ARG002
        self._leave_function()

        if self._is_top_level():
            # Exiting top-level function that's not in a class
            self.current_top_level_name = ""

    def visit_ClassDef(self, node: cst.ClassDef) -> None:
        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._visit_class()

    def leave_ClassDef(self, original_node: cst.ClassDef) -> None:  # noqa: ARG002
        self._leave_class()

        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:
        # Only handle top-level assignments
        if self._is_top_level():
            for target in node.targets:
                # Extract all variable names from the target
                names = extract_names_from_targets(target.target)

                # Check if any of these names are top-level definitions we're tracking
                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 the first tracked name as the current top-level name (for dependency tracking)
                    self.current_top_level_name = tracked_names[0]

    def leave_Assign(self, original_node: cst.Assign) -> None:  # noqa: ARG002
        if self.processing_variable:
            self.processing_variable = False
            self.current_variable_names.clear()
            self.current_top_level_name = ""

    def visit_Name(self, node: cst.Name) -> None:
        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

        # Check if name is a top-level definition we're tracking
        if name in self.definitions and name != self.current_top_level_name:
            self.definitions[self.current_top_level_name].dependencies.add(name)


class TestUsageMarker:
    """Marks definitions that are used by tests (separate from collecting dependencies)."""

    def __init__(self, definitions: dict[str, UsageInfo]) -> None:
        self.definitions = definitions

    def mark_used_definitions(self) -> None:
        """Find all test functions and mark them and their dependencies as used."""
        # First identify all test functions
        test_functions = [name for name, info in self.definitions.items() if info.used_by_test]

        # For each test function, mark all its dependencies as used
        for test_name in test_functions:
            for dep in self.definitions[test_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 by a test recursively."""
        if name not in self.definitions:
            return

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

        self.definitions[name].used_by_test = 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: PLR0911
    node: cst.CSTNode, definitions: dict[str, UsageInfo]
) -> tuple[cst.CSTNode | None, bool]:
    """Recursively filter the node to remove unused definitions.

    Args:
        node: The CST node to process
        definitions: Dictionary of definition info

    Returns:
        (filtered_node, used_by_test):
          filtered_node: The modified CST node or None if it should be removed
          used_by_test: True if this node or any child is used by tests

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

    # Handle function definitions
    if isinstance(node, cst.FunctionDef):
        function_name = node.name.value
        # Test functions are always kept
        if is_test_function_name(function_name):
            return node, True
        # Check if this function is in our definitions and used
        if function_name in definitions and definitions[function_name].used_by_test:
            return node, True
        return None, False

    # Handle class definitions
    if isinstance(node, cst.ClassDef):
        class_name = node.name.value
        # Check if this class is in our definitions and used
        if class_name in definitions and definitions[class_name].used_by_test:
            return node, True
        return None, False

    # Handle assignments (Assign and AnnAssign)
    if isinstance(node, cst.Assign):
        for target in node.targets:
            names = extract_names_from_targets(target.target)
            for name in names:
                if name in definitions and definitions[name].used_by_test:
                    return node, True
        return None, False

    if isinstance(node, cst.AnnAssign | cst.AugAssign):
        names = extract_names_from_targets(node.target)
        for name in names:
            if name in definitions and definitions[name].used_by_test:
                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

    updates = {}
    found_used = False

    for section in section_names:
        original_content = getattr(node, section, None)
        if isinstance(original_content, list | tuple):
            new_children = []
            section_found_used = False

            for child in original_content:
                filtered, used = remove_unused_definitions_recursively(child, definitions)
                if filtered:
                    new_children.append(filtered)
                section_found_used |= used

            if new_children or section_found_used:
                found_used |= section_found_used
                updates[section] = new_children
        elif original_content is not None:
            filtered, used = remove_unused_definitions_recursively(original_content, definitions)
            found_used |= used
            if filtered:
                updates[section] = filtered
    if not found_used:
        return None, False
    if updates:
        return node.with_changes(**updates), found_used

    return node, False


def remove_unused_definitions_from_pytest_file(module: cst.Module) -> cst.Module:
    """Analyze a pytest file and remove top level definitions not used by pytest tests.

    Top level definitions, in this context, are only classes, variables or functions.
    If a class function is used, we keep the entire class.
    """
    # Collect all definitions (top level classes, variables or function)
    definitions = collect_top_level_definitions(module)

    # Collect dependencies between definitions using the visitor pattern
    dependency_collector = DependencyCollector(definitions)
    module.visit(dependency_collector)

    # Mark definitions used by tests, and their dependencies recursively
    usage_marker = TestUsageMarker(definitions)
    usage_marker.mark_used_definitions()

    # Apply the recursive removal transformation
    modified_module, _ = remove_unused_definitions_recursively(module, definitions)

    if not isinstance(modified_module, cst.Module):
        return module
    return modified_module


def print_definitions(definitions: dict[str, UsageInfo]) -> None:
    """Print information about each definition without the complex node object, used for debugging."""
    print(f"Found {len(definitions)} definitions:")
    for name, info in sorted(definitions.items()):
        print(f"  - Name: {name}")
        print(f"    Used by test: {info.used_by_test}")
        print(f"    Dependencies: {', '.join(sorted(info.dependencies)) if info.dependencies else 'None'}")
        print()
