"""Import analysis and test file filtering."""

from __future__ import annotations

import ast
import logging
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .models import TestsInFile


log = logging.getLogger(__name__)


class ImportAnalyzer(ast.NodeVisitor):
    """AST-based analyzer that checks whether any qualified names
    from *function_names_to_find* are imported or used in a file.
    """

    def __init__(self, function_names_to_find: set[str]) -> None:
        """Initialize with the set of qualified function names to search for."""
        self.function_names_to_find = function_names_to_find
        self.found_any_target_function: bool = False
        self.found_qualified_name: str | None = None
        self.imported_modules: set[str] = set()
        self.has_dynamic_imports: bool = False
        self.wildcard_modules: set[str] = set()
        self.alias_mapping: dict[str, str] = {}
        self.instance_mapping: dict[str, str] = {}

        self._exact_names = function_names_to_find
        self._prefix_roots: dict[str, list[str]] = {}
        self._dot_names: set[str] = set()
        self._dot_methods: dict[str, set[str]] = {}
        self._class_method_to_target: dict[tuple[str, str], str] = {}

        add_dot_methods = self._dot_methods.setdefault
        add_prefix_roots = self._prefix_roots.setdefault
        dot_names_add = self._dot_names.add
        class_method_to_target = self._class_method_to_target
        for name in function_names_to_find:
            if "." in name:
                root, method = name.rsplit(".", 1)
                dot_names_add(name)
                add_dot_methods(method, set()).add(root)
                class_method_to_target[(root, method)] = name
                root_prefix = name.split(".", 1)[0]
                add_prefix_roots(root_prefix, []).append(name)

    def visit_Import(self, node: ast.Import) -> None:
        """Handle 'import module' statements."""
        if self.found_any_target_function:
            return

        for alias in node.names:
            module_name = alias.asname or alias.name
            self.imported_modules.add(module_name)

            if alias.name == "importlib":
                self.has_dynamic_imports = True

            if module_name in self.function_names_to_find:
                self.found_any_target_function = True
                self.found_qualified_name = module_name
                return
            for target_func in self.function_names_to_find:
                if target_func.startswith(f"{module_name}."):
                    self.found_any_target_function = True
                    self.found_qualified_name = target_func
                    return

    def visit_Assign(self, node: ast.Assign) -> None:
        """Track variable assignments."""
        if self.found_any_target_function:
            return

        value = node.value
        if isinstance(value, ast.Call) and isinstance(value.func, ast.Name):
            class_name = value.func.id
            if class_name in self.imported_modules:
                original_class = self.alias_mapping.get(class_name, class_name)
                instance_mapping = self.instance_mapping
                for target in node.targets:
                    if isinstance(target, ast.Name):
                        instance_mapping[target.id] = original_class

        self.generic_visit(node)

    def visit_ImportFrom(  # noqa: C901, PLR0912
        self, node: ast.ImportFrom
    ) -> None:
        """Handle 'from module import name' statements."""
        if self.found_any_target_function:
            return

        mod = node.module
        if not mod:
            return

        fnames = self._exact_names
        proots = self._prefix_roots

        for alias in node.names:
            aname = alias.name
            if aname == "*":
                self.wildcard_modules.add(mod)
                continue

            imported_name = alias.asname or aname
            self.imported_modules.add(imported_name)

            if alias.asname:
                self.alias_mapping[imported_name] = aname

            if mod == "importlib" and aname == "import_module":
                self.has_dynamic_imports = True

            qname = f"{mod}.{aname}"

            if aname in fnames:
                self.found_any_target_function = True
                self.found_qualified_name = aname
                return
            if qname in fnames:
                self.found_any_target_function = True
                self.found_qualified_name = qname
                return

            for target_func in fnames:
                if "." in target_func:
                    class_name, _method = target_func.split(".", 1)
                    if aname == class_name and not alias.asname:
                        self.found_any_target_function = True
                        self.found_qualified_name = target_func
                        return

            prefix = qname + "."
            candidates = proots.get(qname, ())
            for target_func in candidates:
                if target_func.startswith(prefix):
                    self.found_any_target_function = True
                    self.found_qualified_name = target_func
                    return

    def visit_Attribute(  # noqa: C901
        self, node: ast.Attribute
    ) -> None:
        """Handle attribute access like obj.func."""
        if self.found_any_target_function:
            return

        node_value = node.value
        node_attr = node.attr

        val_id = getattr(node_value, "id", None)
        if val_id is not None and val_id in self.imported_modules:
            if node_attr in self.function_names_to_find:
                self.found_any_target_function = True
                self.found_qualified_name = node_attr
                return
            roots_possible = self._dot_methods.get(node_attr)
            if roots_possible:
                imported_name = val_id
                original_name = self.alias_mapping.get(
                    imported_name, imported_name
                )
                if original_name in roots_possible:
                    self.found_any_target_function = True
                    self.found_qualified_name = self._class_method_to_target[
                        (original_name, node_attr)
                    ]
                    return
                if imported_name in roots_possible:
                    self.found_any_target_function = True
                    self.found_qualified_name = (
                        self._class_method_to_target.get(
                            (imported_name, node_attr),
                            f"{imported_name}.{node_attr}",
                        )
                    )
                    return

        if val_id is not None and val_id in self.instance_mapping:
            class_name = self.instance_mapping[val_id]
            roots_possible = self._dot_methods.get(node_attr)
            if roots_possible and class_name in roots_possible:
                self.found_any_target_function = True
                self.found_qualified_name = self._class_method_to_target[
                    (class_name, node_attr)
                ]
                return

        if (
            self.has_dynamic_imports
            and node_attr in self.function_names_to_find
        ):
            self.found_any_target_function = True
            self.found_qualified_name = node_attr
            return

        if not self.found_any_target_function:
            ast.NodeVisitor.generic_visit(self, node)

    def visit_Call(self, node: ast.Call) -> None:
        """Handle function calls, particularly __import__."""
        if self.found_any_target_function:
            return

        if isinstance(node.func, ast.Name) and node.func.id == "__import__":
            self.has_dynamic_imports = True

        self.generic_visit(node)

    def visit_Name(self, node: ast.Name) -> None:
        """Handle direct name usage."""
        if self.found_any_target_function:
            return

        if node.id == "__import__":
            self.has_dynamic_imports = True

        if node.id in self.function_names_to_find:
            self.found_any_target_function = True
            self.found_qualified_name = node.id
            return

        for wmod in self.wildcard_modules:
            for target_func in self.function_names_to_find:
                if target_func.startswith(f"{wmod}.") and target_func.endswith(
                    f".{node.id}"
                ):
                    self.found_any_target_function = True
                    self.found_qualified_name = target_func
                    return

        self.generic_visit(node)

    def generic_visit(self, node: ast.AST) -> None:
        """Stop traversal when a target is found."""
        if self.found_any_target_function:
            return
        self._fast_generic_visit(node)

    def _fast_generic_visit(  # noqa: C901, PLR0912
        self, node: ast.AST
    ) -> None:
        """Iterative traversal avoiding method resolution overhead."""
        if self.found_any_target_function:
            return

        visit_cache = type(self).__dict__
        node_fields = node._fields
        stack: list[tuple[tuple[str, ...], ast.AST]] = [
            (node_fields, node),
        ]
        append = stack.append
        pop = stack.pop

        while stack:
            fields, curr_node = pop()
            for field in fields:
                value = getattr(curr_node, field, None)
                if isinstance(value, list):
                    for item in value:
                        if self.found_any_target_function:
                            return
                        if isinstance(item, ast.AST):
                            cls = item.__class__.__name__
                            meth = visit_cache.get("visit_" + cls)
                            if meth is not None:
                                meth(self, item)
                            else:
                                append((item._fields, item))
                    continue
                if isinstance(value, ast.AST):
                    if self.found_any_target_function:
                        return
                    cls = value.__class__.__name__
                    meth = visit_cache.get("visit_" + cls)
                    if meth is not None:
                        meth(self, value)
                    else:
                        append((value._fields, value))


def analyze_imports_in_test_file(
    test_file_path: Path | str,
    target_functions: set[str],
) -> bool:
    """Analyze a test file to see if it imports any target functions."""
    try:
        with Path(test_file_path).open("r", encoding="utf-8") as f:
            source_code = f.read()
        tree = ast.parse(source_code, filename=str(test_file_path))
        analyzer = ImportAnalyzer(target_functions)
        analyzer.visit(tree)
    except (SyntaxError, FileNotFoundError):
        return True

    if analyzer.found_any_target_function:
        return True

    # Be conservative with dynamic imports
    if analyzer.has_dynamic_imports:
        for target_func in target_functions:
            if target_func in source_code:
                return True

    return False


def filter_test_files_by_imports(
    file_to_test_map: dict[Path, list[TestsInFile]],
    target_functions: set[str],
) -> dict[Path, list[TestsInFile]]:
    """Filter test files based on import analysis.

    Returns a subset of *file_to_test_map* containing only
    files that reference *target_functions*.
    """
    if not target_functions:
        return file_to_test_map

    filtered_map = {
        test_file: test_fns
        for test_file, test_fns in file_to_test_map.items()
        if analyze_imports_in_test_file(test_file, target_functions)
    }

    log.debug(
        "analyzed %d test files for imports, "
        "filtered down to %d relevant files",
        len(file_to_test_map),
        len(filtered_map),
    )
    return filtered_map


def glob_test_files(test_root: Path) -> list[Path]:
    """Glob for test files recursively under *test_root*.

    Finds files matching ``test_*.py`` and ``*_test.py``.
    """
    test_prefix = set(test_root.rglob("test_*.py"))
    test_suffix = set(test_root.rglob("*_test.py"))
    return sorted(test_prefix | test_suffix)
