"""Discover optimizable functions in Python source files."""

from __future__ import annotations

import ast
import logging
import os
import random
import warnings
from pathlib import Path
from typing import Any

import attrs

from .._model import FunctionParent, FunctionToOptimize
from ..test_discovery.linking import module_name_from_file_path
from ._reference_graph import path_belongs_to_site_packages

log = logging.getLogger(__name__)


def has_return_statement(
    func_node: ast.FunctionDef | ast.AsyncFunctionDef,
) -> bool:
    """Check if function contains a return statement (recursive)."""
    return any(isinstance(node, ast.Return) for node in ast.walk(func_node))


def is_pytest_fixture(
    func_node: ast.FunctionDef | ast.AsyncFunctionDef,
) -> bool:
    """Check if the function is a pytest fixture."""
    for dec in func_node.decorator_list:
        # @pytest.fixture or @fixture
        unwrapped = dec.func if isinstance(dec, ast.Call) else dec
        if isinstance(unwrapped, ast.Attribute):
            # pytest.fixture
            if (
                unwrapped.attr == "fixture"
                and isinstance(unwrapped.value, ast.Name)
                and unwrapped.value.id == "pytest"
            ):
                return True
        elif isinstance(unwrapped, ast.Name) and unwrapped.id == "fixture":
            return True
    return False


def is_property(
    func_node: ast.FunctionDef | ast.AsyncFunctionDef,
) -> bool:
    """Check if the function is a property."""
    for dec in func_node.decorator_list:
        if isinstance(dec, ast.Name) and dec.id in (
            "property",
            "cached_property",
        ):
            return True
    return False


def discover_functions(
    source: str,
    file_path: Path,
) -> list[FunctionToOptimize]:
    """Discover optimizable functions in Python *source* code."""
    try:
        tree = ast.parse(source, filename=str(file_path))
    except SyntaxError:
        return []

    functions: list[FunctionToOptimize] = []
    # Stack of (node, class_chain) where class_chain tracks ClassDef ancestors
    stack: list[tuple[ast.AST, list[ast.ClassDef]]] = [(tree, [])]

    while stack:
        node, class_chain = stack.pop()

        # Process ClassDef: push children with extended chain
        if isinstance(node, ast.ClassDef):
            new_chain = [*class_chain, node]
            stack.extend(
                (child, new_chain) for child in ast.iter_child_nodes(node)
            )
            continue

        # Process FunctionDef/AsyncFunctionDef at module or class level
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            # Skip if no return statement
            if not has_return_statement(node):
                continue

            # Skip pytest fixtures
            if is_pytest_fixture(node):
                continue

            # Skip properties
            if is_property(node):
                continue

            # Skip nested functions (functions inside functions).
            # Functions can only be at module level or in classes.
            # If we got here, we're at the right level.

            # Build parents tuple
            parents = tuple(
                FunctionParent(cls.name, cls.__class__.__name__)
                for cls in class_chain
            )

            is_method = bool(class_chain)
            is_async = isinstance(node, ast.AsyncFunctionDef)

            functions.append(
                FunctionToOptimize(
                    function_name=node.name,
                    file_path=file_path,
                    parents=parents,
                    starting_line=node.lineno,
                    ending_line=node.end_lineno or node.lineno,
                    starting_col=node.col_offset,
                    ending_col=node.end_col_offset or node.col_offset,
                    is_async=is_async,
                    is_method=is_method,
                )
            )
            # Don't push function children to avoid nested functions

        # For other nodes at module level, push children
        elif not class_chain:
            stack.extend(
                (child, class_chain) for child in ast.iter_child_nodes(node)
            )

    return functions


@attrs.frozen
class FunctionProperties:
    """Properties of a function discovered by AST inspection."""

    is_top_level: bool
    has_args: bool | None
    is_staticmethod: bool | None
    is_classmethod: bool | None
    staticmethod_class_name: str | None


class TopLevelFunctionOrMethodVisitor(ast.NodeVisitor):
    """Visit an AST to check if a named function/method is top-level."""

    def __init__(
        self,
        file_name: Path,
        function_or_method_name: str,
        class_name: str | None = None,
        line_no: int | None = None,
    ) -> None:
        """Store lookup parameters."""
        self.file_name = file_name
        self.class_name = class_name
        self.function_name = function_or_method_name
        self.is_top_level = False
        self.function_has_args: bool | None = None
        self.line_no = line_no
        self.is_staticmethod = False
        self.is_classmethod = False

    def visit_FunctionDef(
        self,
        node: ast.FunctionDef,
    ) -> None:
        """Check top-level function definitions."""
        if self.class_name is None and node.name == self.function_name:
            self.is_top_level = True
            self.function_has_args = any(
                (
                    bool(node.args.args),
                    bool(node.args.kwonlyargs),
                    bool(node.args.kwarg),
                    bool(node.args.posonlyargs),
                    bool(node.args.vararg),
                )
            )

    def visit_AsyncFunctionDef(
        self,
        node: ast.AsyncFunctionDef,
    ) -> None:
        """Check top-level async function definitions."""
        if self.class_name is None and node.name == self.function_name:
            self.is_top_level = True
            self.function_has_args = any(
                (
                    bool(node.args.args),
                    bool(node.args.kwonlyargs),
                    bool(node.args.kwarg),
                    bool(node.args.posonlyargs),
                    bool(node.args.vararg),
                )
            )

    def visit_ClassDef(
        self,
        node: ast.ClassDef,
    ) -> None:
        """Check class methods and static methods."""
        if node.name == self.class_name:
            for body_node in node.body:
                if (
                    isinstance(
                        body_node,
                        (ast.FunctionDef, ast.AsyncFunctionDef),
                    )
                    and body_node.name == self.function_name
                ):
                    self.is_top_level = True
                    if any(
                        isinstance(d, ast.Name) and d.id == "classmethod"
                        for d in body_node.decorator_list
                    ):
                        self.is_classmethod = True
                    elif any(
                        isinstance(d, ast.Name) and d.id == "staticmethod"
                        for d in body_node.decorator_list
                    ):
                        self.is_staticmethod = True
                    return
        elif self.line_no:
            for body_node in node.body:
                if (
                    isinstance(
                        body_node,
                        (ast.FunctionDef, ast.AsyncFunctionDef),
                    )
                    and body_node.name == self.function_name
                    and body_node.lineno in {self.line_no, self.line_no + 1}
                    and any(
                        isinstance(d, ast.Name) and d.id == "staticmethod"
                        for d in body_node.decorator_list
                    )
                ):
                    self.is_staticmethod = True
                    self.is_top_level = True
                    self.class_name = node.name
                    return


def inspect_top_level_functions_or_methods(
    file_name: Path,
    function_or_method_name: str,
    class_name: str | None = None,
    line_no: int | None = None,
) -> FunctionProperties | None:
    """Inspect whether a function/method is top-level in *file_name*."""
    try:
        source = file_name.read_text(encoding="utf-8")
    except OSError:
        log.warning("Failed to read %s", file_name)
        return None
    try:
        ast_module = ast.parse(source)
    except SyntaxError:
        log.debug("Failed to parse %s", file_name)
        return None
    visitor = TopLevelFunctionOrMethodVisitor(
        file_name=file_name,
        function_or_method_name=function_or_method_name,
        class_name=class_name,
        line_no=line_no,
    )
    visitor.visit(ast_module)
    staticmethod_class_name = (
        visitor.class_name if visitor.is_staticmethod else None
    )
    return FunctionProperties(
        is_top_level=visitor.is_top_level,
        has_args=visitor.function_has_args,
        is_staticmethod=visitor.is_staticmethod,
        is_classmethod=visitor.is_classmethod,
        staticmethod_class_name=staticmethod_class_name,
    )


_VCS_EXCLUDES = frozenset({".git", ".hg", ".svn"})

_PYTHON_DIR_EXCLUDES = frozenset(
    {
        "__pycache__",
        ".venv",
        "venv",
        ".tox",
        ".nox",
        ".eggs",
        ".mypy_cache",
        ".ruff_cache",
        ".pytest_cache",
        ".hypothesis",
        "htmlcov",
        ".pytype",
        ".pyre",
        ".pybuilder",
        ".ipynb_checkpoints",
        ".codeflash",
        ".cache",
        ".complexipy_cache",
        "build",
        "dist",
        "sdist",
    }
)

_ALL_DIR_EXCLUDES = _VCS_EXCLUDES | _PYTHON_DIR_EXCLUDES

_SUFFIX_DIR_EXCLUDES = (".egg-info",)
_PREFIX_DIR_EXCLUDES = (".coverage", ".pyright")


def find_all_functions_in_file(
    file_path: Path,
) -> dict[Path, list[FunctionToOptimize]]:
    """Find all optimizable functions in a Python file."""
    try:
        source = file_path.read_text(encoding="utf-8")
    except OSError:
        log.warning("Failed to read %s", file_path)
        return {}
    try:
        ast.parse(source, filename=str(file_path))
    except SyntaxError:
        log.debug("Failed to parse %s", file_path)
        return {}
    fns = discover_functions(source, file_path)
    fns.sort(key=lambda f: f.starting_line or 0)
    return {file_path: fns}


def get_python_files(
    module_root_path: Path,
    ignore_paths: list[Path] | None = None,
) -> list[Path]:
    """Walk *module_root_path* and return all ``.py`` file paths."""
    if ignore_paths is None:
        ignore_paths = []

    ignore_dirs: set[str] = set()
    ignore_files: set[Path] = set()
    for p in ignore_paths:
        p = Path(p) if not isinstance(p, Path) else p  # noqa: PLW2901
        if p.is_file():
            ignore_files.add(p)
        else:
            ignore_dirs.add(str(p))

    files: list[Path] = []
    for dirpath, dirnames, filenames in os.walk(module_root_path):
        dirnames[:] = [
            d
            for d in dirnames
            if d not in _ALL_DIR_EXCLUDES
            and not d.endswith(_SUFFIX_DIR_EXCLUDES)
            and not d.startswith(_PREFIX_DIR_EXCLUDES)
            and str(Path(dirpath) / d) not in ignore_dirs
        ]
        for fname in filenames:
            if fname.endswith(".py"):
                fpath = Path(dirpath, fname)
                if fpath not in ignore_files:
                    files.append(fpath)
    return files


def get_all_files_and_functions(
    module_root_path: Path,
    ignore_paths: list[Path],
) -> dict[Path, list[FunctionToOptimize]]:
    """Discover all functions in ``.py`` files under *module_root_path*."""
    functions: dict[Path, list[FunctionToOptimize]] = {}
    for file_path in get_python_files(module_root_path, ignore_paths):
        functions.update(find_all_functions_in_file(file_path).items())
    files_list = list(functions.items())
    random.shuffle(files_list)
    return dict(files_list)


def get_blocklisted_functions() -> dict[str, set[str]]:
    """Return blocklisted functions from the platform API.

    Stub — returns empty dict until the platform client is wired.
    """
    return {}


def filter_functions(  # noqa: C901, PLR0912, PLR0913, PLR0915
    modified_functions: dict[Path, list[FunctionToOptimize]],
    tests_root: Path,
    ignore_paths: list[Path],
    project_root: Path,
    module_root: Path,
    previous_checkpoint_functions: (dict[str, dict[str, Any]] | None) = None,
    *,
    disable_logs: bool = False,
) -> tuple[dict[Path, list[FunctionToOptimize]], int]:
    """Filter discovered functions, removing tests and non-optimizable."""
    from ..benchmarking._file_filtering import (  # noqa: PLC0415
        ignored_submodule_paths,
    )

    resolved_project_root = project_root.resolve()
    filtered: dict[Path, list[FunctionToOptimize]] = {}
    blocklist_funcs = get_blocklisted_functions()

    submodule_paths = ignored_submodule_paths(module_root)

    functions_count: int = 0
    test_functions_removed: int = 0
    non_modules_removed: int = 0
    site_packages_removed: int = 0
    ignore_paths_removed: int = 0
    malformed_paths: int = 0
    submodule_ignored: int = 0
    blocklist_removed: int = 0
    checkpoint_removed: int = 0

    tests_root_str = os.path.normcase(str(tests_root))
    module_root_str = os.path.normcase(str(module_root))
    project_root_str = os.path.normcase(str(project_root))

    tests_root_overlaps_source = tests_root_str in (
        module_root_str,
        project_root_str,
    ) or module_root_str.startswith(tests_root_str + os.sep)

    test_file_name_patterns = (
        ".test.",
        ".spec.",
        "_test.",
        "_spec.",
    )
    test_dir_patterns = (
        os.sep + "test" + os.sep,
        os.sep + "tests" + os.sep,
        os.sep + "__tests__" + os.sep,
    )

    def is_test_file(file_path_normalized: str) -> bool:
        if tests_root_overlaps_source:
            file_lower = file_path_normalized.lower()
            basename = Path(file_lower).name
            if basename.startswith("test_") or basename == "conftest.py":
                return True
            if any(p in file_lower for p in test_file_name_patterns):
                return True
            if project_root_str and file_lower.startswith(
                project_root_str.lower()
            ):
                relative = file_lower[len(project_root_str) :]
                return any(p in relative for p in test_dir_patterns)
            return False
        return file_path_normalized.startswith(tests_root_str + os.sep)

    for file_path_path, functions in modified_functions.items():
        fns = functions
        file_path = str(file_path_path)
        file_path_normalized = os.path.normcase(file_path)
        if is_test_file(file_path_normalized):
            test_functions_removed += len(fns)
            continue
        if file_path_path in ignore_paths or any(
            file_path_normalized.startswith(os.path.normcase(str(ip)) + os.sep)
            for ip in ignore_paths
        ):
            ignore_paths_removed += 1
            continue
        if file_path_path in submodule_paths or any(
            file_path_normalized.startswith(os.path.normcase(str(sp)) + os.sep)
            for sp in submodule_paths
        ):
            submodule_ignored += 1
            continue
        if path_belongs_to_site_packages(Path(file_path)):
            site_packages_removed += len(fns)
            continue
        if not file_path_normalized.startswith(module_root_str + os.sep):
            non_modules_removed += len(fns)
            continue

        try:
            ast.parse(
                f"import {module_name_from_file_path(Path(file_path), resolved_project_root)}"
            )
        except (SyntaxError, ValueError):
            malformed_paths += 1
            continue

        if blocklist_funcs:
            tmp = []
            for fn in fns:
                if (
                    fn.file_path.name in blocklist_funcs
                    and fn.qualified_name in blocklist_funcs[fn.file_path.name]
                ):
                    blocklist_removed += 1
                    continue
                tmp.append(fn)
            fns = tmp

        if previous_checkpoint_functions:
            tmp = []
            for fn in fns:
                qn = fn.qualified_name_with_modules_from_root(
                    resolved_project_root
                )
                if qn in previous_checkpoint_functions:
                    checkpoint_removed += 1
                    continue
                tmp.append(fn)
            fns = tmp

        filtered[file_path_path] = fns
        functions_count += len(fns)

    if not disable_logs:
        info = {
            "Test functions removed": test_functions_removed,
            "Site-package functions removed": site_packages_removed,
            "Non-importable file paths": malformed_paths,
            "Functions outside module-root": non_modules_removed,
            "Files from ignored paths": ignore_paths_removed,
            "Files from ignored submodules": submodule_ignored,
            "Blocklisted functions removed": blocklist_removed,
            "Functions skipped from checkpoint": checkpoint_removed,
        }
        parts = [
            f"{label}: {count}" for label, count in info.items() if count > 0
        ]
        if parts:
            log.debug("Ignored functions: %s", "; ".join(parts))
    return (
        {k: v for k, v in filtered.items() if v},
        functions_count,
    )


def get_functions_to_optimize(  # noqa: PLR0913
    optimize_all: str | None,
    replay_test: list[Path] | None,
    file: Path | str | None,
    only_get_this_function: str | None,
    test_cfg: Any,
    ignore_paths: list[Path],
    project_root: Path,
    module_root: Path,
    previous_checkpoint_functions: (dict[str, dict[str, str]] | None) = None,
) -> tuple[dict[Path, list[FunctionToOptimize]], int, Path | None]:
    """Discover and filter functions to optimize.

    Returns ``(functions_dict, count, trace_file_path)``.
    """
    functions: dict[Path, list[FunctionToOptimize]]
    trace_file_path: Path | None = None

    with warnings.catch_warnings():
        warnings.simplefilter(action="ignore", category=SyntaxWarning)
        if optimize_all:
            functions = get_all_files_and_functions(
                Path(optimize_all), ignore_paths
            )
        elif file is not None:
            file = Path(file) if isinstance(file, str) else file
            functions = find_all_functions_in_file(file)
            if only_get_this_function is not None:
                split_fn = only_get_this_function.split(".")
                if len(split_fn) == 2:  # noqa: PLR2004
                    class_name, fn_name = split_fn
                else:
                    class_name = None
                    fn_name = split_fn[0]
                found = None
                for fn in functions.get(file, []):
                    if fn_name == fn.function_name and (
                        class_name is None
                        or class_name == fn.top_level_parent_name
                    ):
                        found = fn
                if found is not None:
                    functions[file] = [found]
                else:
                    functions = {}
        else:
            functions = {}

    tests_root = getattr(test_cfg, "tests_root", "tests")
    filtered, count = filter_functions(
        functions,
        Path(tests_root),
        ignore_paths,
        project_root,
        module_root,
        previous_checkpoint_functions,
    )
    return filtered, count, trace_file_path


def _is_test_file_by_pattern(file_path: Path) -> bool:
    """Check if a file is a test file using naming conventions.

    Used when tests_root overlaps with module_root, so directory-based
    filtering would incorrectly exclude all source files. Falls back to
    filename and directory patterns.
    """
    name = file_path.name.lower()
    if name.startswith("test_") or name == "conftest.py":
        return True
    test_name_patterns = (".test.", ".spec.", "_test.", "_spec.")
    if any(p in name for p in test_name_patterns):
        return True
    path_str = str(file_path).lower()
    test_dir_patterns = (
        os.sep + "test" + os.sep,
        os.sep + "tests" + os.sep,
        os.sep + "__tests__" + os.sep,
    )
    return any(p in path_str for p in test_dir_patterns)


def filter_files_optimized(
    file_path: Path,
    tests_root: Path,
    ignore_paths: list[Path],
    module_root: Path,
) -> bool:
    """Return True if *file_path* should be considered for optimization."""
    from ..benchmarking._file_filtering import (  # noqa: PLC0415
        ignored_submodule_paths,
    )
    from ._reference_graph import (  # noqa: PLC0415
        path_belongs_to_site_packages,
    )

    tests_root_overlaps = (
        tests_root == module_root or module_root.is_relative_to(tests_root)
    )
    if tests_root_overlaps:
        if _is_test_file_by_pattern(file_path):
            return False
    elif file_path.is_relative_to(tests_root):
        return False
    if file_path in ignore_paths or any(
        file_path.is_relative_to(p) for p in ignore_paths
    ):
        return False
    if path_belongs_to_site_packages(file_path):
        return False
    if not file_path.is_relative_to(module_root):
        return False
    submodule_paths = ignored_submodule_paths(module_root)
    return not (
        file_path in submodule_paths
        or any(file_path.is_relative_to(sp) for sp in submodule_paths)
    )
