"""Helper discovery via Jedi."""

from __future__ import annotations

import logging
import os
from typing import TYPE_CHECKING, Any

import jedi  # type: ignore[import-untyped]

from .._model import FunctionSource
from ..analysis._reference_graph import path_belongs_to_site_packages
from .resolve import get_jedi_project, get_qualified_name

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionToOptimize

log = logging.getLogger(__name__)


def is_project_path(
    module_path: Path,
    project_root: Path,
) -> bool:
    """Return *True* if *module_path* is under *project_root*."""
    return str(module_path.resolve()).startswith(
        str(project_root.resolve()) + os.sep,
    )


def belongs_to_function(
    name: Any,
    qualified_function_name: str,
) -> bool:
    """Return *True* if *name* is defined inside *qualified_function_name*."""
    try:
        if (
            name.full_name
            and name.full_name.startswith(name.module_name)
            and get_qualified_name(
                name.module_name,
                name.full_name,
            )
            == qualified_function_name
        ):
            return False  # name IS the function, not internal
        parent = name.parent()
        if parent is not None and parent.type == "function":
            return (
                parent.full_name is not None
                and parent.full_name.startswith(
                    parent.module_name,
                )
                and get_qualified_name(
                    parent.module_name,
                    parent.full_name,
                )
                == qualified_function_name
            )
    except (ValueError, AttributeError):
        pass
    return False


def group_refs_by_parent(
    file_path: Path,
    project: Any,
) -> dict[str, list[Any]]:
    """Group references in *file_path* by parent function."""
    script = jedi.Script(path=file_path, project=project)
    refs = script.get_names(
        all_scopes=True,
        definitions=False,
        references=True,
    )

    result: dict[str, list[Any]] = {}
    for ref in refs:
        try:
            parent = ref.parent()
            if parent is None or parent.type != "function":
                continue
            if not parent.full_name or not parent.full_name.startswith(
                parent.module_name,
            ):
                continue
            parent_qn = get_qualified_name(
                parent.module_name,
                parent.full_name,
            )
            # Skip self-references (recursion)
            if ref.full_name and ref.full_name.startswith(ref.module_name):
                ref_qn = get_qualified_name(
                    ref.module_name,
                    ref.full_name,
                )
                if ref_qn == parent_qn:
                    continue
            result.setdefault(parent_qn, []).append(ref)
        except (ValueError, AttributeError):
            continue

    return result


MAX_QUALIFIED_DEPTH = 2


def is_valid_helper_defn(
    defn: Any,
    qualified_function_name: str,
    project_root: Path,
) -> bool:
    """Return *True* if *defn* is a valid in-project helper."""
    defn_path = defn.module_path
    if defn_path is None or not defn.full_name:
        return False
    if path_belongs_to_site_packages(defn_path):
        return False
    if not is_project_path(defn_path, project_root):
        return False
    if belongs_to_function(defn, qualified_function_name):
        return False
    return defn.full_name.startswith(defn.module_name) and defn.type in (
        "function",
        "class",
        "statement",
    )


def resolve_ref_to_helper(
    ref: Any,
    qualified_function_name: str,
    project_root: Path,
) -> FunctionSource | None:
    """Resolve a Jedi reference to a *FunctionSource*, or *None*."""
    defns = ref.goto(
        follow_imports=True,
        follow_builtin_imports=False,
    )
    if not defns:
        return None
    defn = defns[0]

    if not is_valid_helper_defn(
        defn,
        qualified_function_name,
        project_root,
    ):
        return None

    defn_qn = get_qualified_name(
        defn.module_name,
        defn.full_name,
    )
    # Skip self-references (recursion)
    if defn_qn == qualified_function_name:
        return None
    if defn.type == "class":
        fqn = f"{defn.full_name}.__init__"
        func_name = "__init__"
        defn_qn = f"{defn_qn}.__init__"
    else:
        fqn = defn.full_name
        func_name = defn.name
    if len(defn_qn.split(".")) > MAX_QUALIFIED_DEPTH:
        return None

    # Normalize path to canonical form under project root
    file_path = defn.module_path
    try:
        rel = file_path.resolve().relative_to(project_root.resolve())
        file_path = project_root / rel
    except ValueError:
        pass

    return FunctionSource(
        file_path=file_path,
        qualified_name=defn_qn,
        fully_qualified_name=fqn,
        source_code=defn.get_line_code(),
        only_function_name=func_name,
        definition_type=defn.type,
    )


def discover_helpers_by_names(
    file_paths_to_names: dict[Path, set[str]],
    project_root: Path,
    refs_cache: dict[Path, dict[str, list[Any]]] | None = None,
) -> dict[Path, set[FunctionSource]]:
    """One-level helper discovery for functions identified by name.

    *refs_cache*, when provided, is checked and populated in-place
    to avoid redundant Jedi ``get_names`` calls on files that
    appear in multiple discovery rounds.
    """
    project = get_jedi_project(str(project_root))
    result: dict[Path, set[FunctionSource]] = {}

    for file_path, qualified_names in file_paths_to_names.items():
        if refs_cache is not None and file_path in refs_cache:
            refs_by_parent = refs_cache[file_path]
        else:
            refs_by_parent = group_refs_by_parent(
                file_path,
                project,
            )
            if refs_cache is not None:
                refs_cache[file_path] = refs_by_parent

        for qn in qualified_names:
            for ref in refs_by_parent.get(qn, []):
                try:
                    fs = resolve_ref_to_helper(
                        ref,
                        qn,
                        project_root,
                    )
                    if fs is not None:
                        result.setdefault(
                            fs.file_path,
                            set(),
                        ).add(fs)
                except Exception:  # noqa: PERF203
                    log.exception(
                        "Error resolving ref in %s",
                        file_path,
                    )
                    continue

    return result


def discover_helpers(
    function: FunctionToOptimize,
    project_root: Path,
) -> dict[Path, set[FunctionSource]]:
    """Discover helpers referenced by *function*, two levels deep."""
    # Level 1: direct helpers
    direct = discover_helpers_by_names(
        {function.file_path: {function.qualified_name}},
        project_root,
    )

    # Build level 2 input from discovered helpers
    level2_input: dict[Path, set[str]] = {}
    for sources in direct.values():
        for fs in sources:
            level2_input.setdefault(
                fs.file_path,
                set(),
            ).add(fs.qualified_name)
            # For class methods, also discover __init__
            parts = fs.qualified_name.split(".")
            if len(parts) >= MAX_QUALIFIED_DEPTH:
                level2_input.setdefault(
                    fs.file_path,
                    set(),
                ).add(f"{parts[0]}.__init__")

    # Level 2: transitive helpers
    transitive = discover_helpers_by_names(
        level2_input,
        project_root,
    )

    # Merge both levels
    result = dict(direct)
    for path, sources in transitive.items():
        result[path] = result.get(path, set()) | sources

    return result
