# pytools/core/extract_imports.py

import ast
import sys
from pathlib import Path
from typing import Iterable, NamedTuple


_STDLIB_ROOTS: frozenset[str] = frozenset(sys.stdlib_module_names)


class ExtractedImports(NamedTuple):
    """Import names grouped into third-party, local-relative, stdlib, and ignored."""

    third_party: list[str]
    local: list[str]
    stdlib: list[str]
    ignored: list[tuple[str, list[str]]]


def _relative_import_label(node: ast.ImportFrom) -> str:
    return "." * node.level + (node.module or "")


def extract_imports(
    path: Path,
    ignore_globs: Iterable[str] = (),
    internal_roots: Iterable[str] = (),
) -> ExtractedImports:
    internal = frozenset(internal_roots)
    ignore_patterns = list(ignore_globs)
    third_party_roots: set[str] = set()
    local_imports: set[str] = set()
    stdlib_roots: set[str] = set()
    ignored_by_pattern: dict[str, set[str]] = {pattern: set() for pattern in ignore_patterns}

    for file in path.rglob("*.py"):
        tree = ast.parse(file.read_text(encoding="utf-8"), filename=str(file))

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for n in node.names:
                    import_name = n.name
                    root = import_name.split(".")[0]
                    if any(
                        _record_ignored(import_name, pattern, ignored_by_pattern)
                        for pattern in ignore_patterns
                    ):
                        continue
                    if root in internal:
                        continue
                    if root in _STDLIB_ROOTS:
                        stdlib_roots.add(root)
                    else:
                        third_party_roots.add(root)

            elif isinstance(node, ast.ImportFrom):
                if node.level > 0:
                    local_label = _relative_import_label(node)
                    if any(
                        _record_ignored(local_label, pattern, ignored_by_pattern)
                        for pattern in ignore_patterns
                    ):
                        continue
                    local_imports.add(local_label)
                elif node.module:
                    import_name = node.module
                    root = import_name.split(".")[0]
                    if any(
                        _record_ignored(import_name, pattern, ignored_by_pattern)
                        for pattern in ignore_patterns
                    ):
                        continue
                    if root in internal:
                        continue
                    if root in _STDLIB_ROOTS:
                        stdlib_roots.add(root)
                    else:
                        third_party_roots.add(root)

    ignored = [
        (pattern, sorted(ignored_by_pattern[pattern]))
        for pattern in ignore_patterns
        if ignored_by_pattern[pattern]
    ]
    return ExtractedImports(
        third_party=sorted(third_party_roots),
        local=sorted(local_imports),
        stdlib=sorted(stdlib_roots),
        ignored=ignored,
    )


def _record_ignored(name: str, pattern: str, ignored_by_pattern: dict[str, set[str]]) -> bool:
    if _matches_ignore(name, pattern):
        ignored_by_pattern[pattern].add(name)
        return True
    return False


def _matches_ignore(name: str, pattern: str) -> bool:
    # Support exact ("pkg.mod"), prefix-with-boundary ("pkg.mod" -> "pkg.mod.*"), and fnmatch globs.
    if pattern == name or name.startswith(f"{pattern}."):
        return True
    if any(char in pattern for char in "*?[]"):
        from fnmatch import fnmatch

        return fnmatch(name, pattern)
    return False
