"""Import gathering, scheduling, and transformation for codegen."""

from __future__ import annotations

import ast
import logging
from typing import TYPE_CHECKING

import libcst as cst
import libcst.matchers as m
from libcst.codemod import CodemodContext
from libcst.codemod.visitors import (
    AddImportsVisitor,
    GatherImportsVisitor,
    RemoveImportsVisitor,
)
from libcst.helpers import calculate_module_and_package

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionSource

log = logging.getLogger(__name__)

_SENTINEL = object()


class DottedImportCollector(cst.CSTVisitor):
    """Collect top-level imports as dotted strings.

    ``from pathlib import Path`` becomes ``'pathlib.Path'``.
    """

    def __init__(self) -> None:
        """Initialize with an empty set of collected imports."""
        self.imports: set[str] = set()

    def get_full_dotted_name(
        self,
        expr: cst.BaseExpression,
    ) -> str:
        """Return the dotted form of *expr*."""
        if isinstance(expr, cst.Name):
            return expr.value
        if isinstance(expr, cst.Attribute):
            return f"{self.get_full_dotted_name(expr.value)}.{expr.attr.value}"
        return ""

    def _collect_imports_from_block(
        self,
        block: cst.IndentedBlock | cst.Module,
    ) -> None:
        """Collect imports from a block's top-level statements."""
        for statement in block.body:
            if not isinstance(statement, cst.SimpleStatementLine):
                continue
            for child in statement.body:
                if isinstance(child, cst.Import):
                    self._collect_plain_import(child)
                elif isinstance(child, cst.ImportFrom):
                    self._collect_from_import(child)

    def _collect_plain_import(
        self,
        node: cst.Import,
    ) -> None:
        """Collect dotted names from a plain import."""
        if isinstance(node.names, cst.ImportStar):
            return
        for alias in node.names:
            module = self.get_full_dotted_name(
                alias.name,
            )
            if alias.asname and isinstance(alias.asname.name, cst.Name):
                asname: str | cst.Attribute = alias.asname.name.value
            else:
                asname = alias.name.value  # type: ignore[assignment]
            if isinstance(asname, cst.Attribute):
                self.imports.add(module)
            else:
                self.imports.add(
                    module if module == asname else f"{module}.{asname}",
                )

    def _collect_from_import(
        self,
        node: cst.ImportFrom,
    ) -> None:
        """Collect dotted names from a from-import."""
        if node.module is None:
            return
        module = self.get_full_dotted_name(node.module)
        if isinstance(node.names, cst.ImportStar):
            return
        for alias in node.names:
            if not isinstance(alias, cst.ImportAlias):
                continue
            if not isinstance(alias.name, cst.Name):
                continue
            name = alias.name.value
            if alias.asname and isinstance(alias.asname.name, cst.Name):
                asname = alias.asname.name.value
            else:
                asname = name
            self.imports.add(f"{module}.{asname}")

    def visit_Module(  # noqa: N802
        self, node: cst.Module
    ) -> None:
        """Collect imports from module body."""
        self._collect_imports_from_block(node)

    def visit_FunctionDef(  # noqa: N802
        self,
        node: cst.FunctionDef,
    ) -> bool:
        """Skip function bodies."""
        return False

    def visit_ClassDef(  # noqa: N802
        self,
        node: cst.ClassDef,
    ) -> bool:
        """Skip class bodies."""
        return False

    def visit_If(  # noqa: N802
        self, node: cst.If
    ) -> None:
        """Collect imports inside ``if`` blocks."""
        if isinstance(node.body, cst.IndentedBlock):
            self._collect_imports_from_block(node.body)

    def visit_Try(  # noqa: N802
        self, node: cst.Try
    ) -> None:
        """Collect imports inside ``try`` blocks."""
        if isinstance(node.body, cst.IndentedBlock):
            self._collect_imports_from_block(node.body)


class FutureAliasedImportTransformer(cst.CSTTransformer):
    """Remove aliased ``__future__`` imports.

    ``from __future__ import annotations as a`` is invalid at
    runtime; this transformer strips the alias or removes the
    entire import line when every name is aliased.
    """

    def leave_ImportFrom(  # noqa: N802
        self,
        original_node: cst.ImportFrom,
        updated_node: cst.ImportFrom,
    ) -> (
        cst.BaseSmallStatement
        | cst.FlattenSentinel[cst.BaseSmallStatement]
        | cst.RemovalSentinel
    ):
        """Strip aliased names from ``__future__`` imports."""
        if (
            (mod := updated_node.module)
            and isinstance(mod, (cst.Attribute, cst.Name))
            and hasattr(mod, "value")
            and mod.value == "__future__"
            and not isinstance(updated_node.names, cst.ImportStar)
            and all(
                m.matches(name, m.ImportAlias()) for name in updated_node.names
            )
        ):
            if names := [
                name for name in updated_node.names if name.asname is None
            ]:
                return updated_node.with_changes(
                    names=names,
                )
            return cst.RemoveFromParent()
        return updated_node


def delete_future_aliased_imports(
    module_code: str,
) -> str:
    """Remove aliased ``__future__`` imports from *module_code*."""
    return (
        cst.parse_module(module_code)
        .visit(FutureAliasedImportTransformer())
        .code
    )


def resolve_star_import(
    module_name: str,
    project_root: Path,
) -> set[str]:
    """Resolve ``from X import *`` to the set of exported names.

    Uses ``__all__`` when present, otherwise falls back to all
    public top-level names.
    """
    try:
        return _resolve_star_import_inner(
            module_name,
            project_root,
        )
    except (OSError, SyntaxError):
        log.warning(
            "Error resolving star import for %s",
            module_name,
        )
        return set()


def _resolve_star_import_inner(
    module_name: str,
    project_root: Path,
) -> set[str]:
    """Resolve star imports by reading the module file."""
    module_path = module_name.replace(".", "/")
    possible = [
        project_root / f"{module_path}.py",
        project_root / f"{module_path}/__init__.py",
    ]

    module_file = next(
        (p for p in possible if p.exists()),
        None,
    )
    if module_file is None:
        log.warning(
            "Could not find module file for %s",
            module_name,
        )
        return set()

    tree = ast.parse(
        module_file.read_text(encoding="utf8"),
    )

    all_names = _extract_all_list(tree)
    if all_names is not None:
        return set(all_names)

    return _collect_public_names(tree)


def _extract_all_list(
    tree: ast.Module,
) -> list[str] | None:
    """Extract the __all__ list from a module AST."""
    for node in ast.walk(tree):
        if (
            isinstance(node, ast.Assign)
            and len(node.targets) == 1
            and isinstance(node.targets[0], ast.Name)
            and node.targets[0].id == "__all__"
            and isinstance(node.value, (ast.List, ast.Tuple))
        ):
            return [
                elt.value
                for elt in node.value.elts
                if isinstance(elt, ast.Constant) and isinstance(elt.value, str)
            ]
    return None


def _collect_public_names(tree: ast.Module) -> set[str]:
    """Collect all public top-level names."""
    names: set[str] = set()
    for node in tree.body:
        _collect_name_from_node(node, names)
    return names


def _collect_name_from_node(
    node: ast.stmt,
    names: set[str],
) -> None:
    """Add the public name defined by an AST statement."""
    if isinstance(
        node,
        (
            ast.FunctionDef,
            ast.AsyncFunctionDef,
            ast.ClassDef,
        ),
    ):
        if not node.name.startswith("_"):
            names.add(node.name)
    elif isinstance(node, ast.Assign):
        _collect_assign_names(node, names)
    elif isinstance(node, ast.AnnAssign) and isinstance(
        node.target,
        ast.Name,
    ):
        if not node.target.id.startswith("_"):
            names.add(node.target.id)
    elif isinstance(
        node, (ast.Import, ast.ImportFrom)
    ) and _is_non_star_import(
        node,
    ):
        for alias in node.names:
            name = alias.asname or alias.name
            if not name.startswith("_"):
                names.add(name)


def _collect_assign_names(
    node: ast.Assign,
    names: set[str],
) -> None:
    """Add public variable names from an assignment."""
    for target in node.targets:
        if isinstance(target, ast.Name) and not (target.id.startswith("_")):
            names.add(target.id)


def _is_non_star_import(node: ast.stmt) -> bool:
    """Return True if the node is an import without stars."""
    return isinstance(node, ast.Import) or (
        isinstance(node, ast.ImportFrom)
        and not any(alias.name == "*" for alias in node.names)
    )


def _collect_dst_referenced_names(
    dst_code: str,
) -> tuple[set[str], bool]:
    """Collect all names referenced in *dst_code*.

    Uses :mod:`ast` (not libcst) for speed.  Returns
    *(names, has_imports)* where *has_imports* indicates
    whether the destination already has import statements.
    """
    try:
        tree = ast.parse(dst_code)
    except SyntaxError:
        return set(), False
    names: set[str] = set()
    has_imports = False
    for node in ast.walk(tree):
        if isinstance(node, ast.Name):
            names.add(node.id)
        elif isinstance(
            node,
            ast.Attribute,
        ) and isinstance(node.value, ast.Name):
            names.add(node.value.id)
        elif isinstance(node, (ast.Import, ast.ImportFrom)):
            has_imports = True
        elif isinstance(node, ast.Constant) and isinstance(
            node.value,
            str,
        ):
            try:
                inner = ast.parse(node.value, mode="eval")
                for inner_node in ast.walk(inner):
                    if isinstance(inner_node, ast.Name):
                        names.add(inner_node.id)
            except SyntaxError:
                pass
    return names, has_imports


def add_needed_imports_from_module(  # noqa: C901, PLR0912, PLR0913
    src_module_code: str | cst.Module,
    dst_module_code: str | cst.Module,
    src_path: Path,
    dst_path: Path,
    project_root: Path,
    *,
    helper_functions: list[FunctionSource] | None = None,
    helper_functions_fqn: set[str] | None = None,
    gathered_imports: (GatherImportsVisitor | None | object) = _SENTINEL,
) -> str:
    """Add needed imports from *src* to *dst* module code.

    Returns the transformed destination code as a string.
    """
    if not helper_functions_fqn:
        helper_functions_fqn = {
            f.fully_qualified_name for f in (helper_functions or [])
        }

    if isinstance(dst_module_code, str):
        dst_fallback = dst_module_code
    else:
        dst_fallback = dst_module_code.code.lstrip(
            "\n",
        )

    dst_mp = calculate_module_and_package(
        project_root,
        dst_path,
    )
    dst_context = CodemodContext(
        filename=src_path.name,
        full_module_name=dst_mp.name,
        full_package_name=dst_mp.package,
    )

    gatherer: GatherImportsVisitor | None
    if gathered_imports is _SENTINEL:
        from ..context.imports import (  # noqa: PLC0415
            gather_source_imports,
        )

        gatherer = gather_source_imports(
            src_module_code,
            src_path,
            project_root,
        )
    else:
        gatherer = gathered_imports  # type: ignore[assignment]

    if gatherer is None:
        return dst_fallback

    collector = DottedImportCollector()
    if isinstance(dst_module_code, str):
        try:
            parsed_dst = cst.parse_module(
                dst_module_code,
            )
        except cst.ParserSyntaxError:
            log.exception(
                "Syntax error in destination module",
            )
            return dst_fallback
    else:
        parsed_dst = dst_module_code
    parsed_dst.visit(collector)

    # Pre-filter: collect names referenced in destination
    # code to avoid adding unused imports.  This keeps the
    # intermediate module small so RemoveImportsVisitor's
    # scope analysis is cheap.
    dst_code_str = (
        parsed_dst.code if isinstance(parsed_dst, cst.Module) else dst_fallback
    )
    (
        dst_referenced_names,
        dst_has_imports,
    ) = _collect_dst_referenced_names(dst_code_str)

    try:
        _schedule_module_imports(
            gatherer,
            collector,
            dst_context,
            dst_referenced_names,
        )
        _schedule_object_imports(
            gatherer,
            collector,
            dst_context,
            helper_functions_fqn,
            project_root,
            dst_referenced_names,
        )
    except Exception:
        log.exception("Error scheduling imports")
        return dst_fallback

    _schedule_alias_imports(
        gatherer,
        collector,
        dst_context,
        helper_functions_fqn,
        dst_referenced_names,
    )

    try:
        transformed = parsed_dst
        if dst_context.scratch.get(
            "AddImportsVisitor",
        ):
            transformed = AddImportsVisitor(
                dst_context,
            ).transform_module(transformed)
        # Skip RemoveImportsVisitor when dst had no
        # pre-existing imports -- the only imports are
        # those just added, which are already
        # pre-filtered to names referenced in dst.
        if dst_has_imports and dst_context.scratch.get(
            "RemoveImportsVisitor",
        ):
            transformed = RemoveImportsVisitor(
                dst_context,
            ).transform_module(transformed)
        return transformed.code.lstrip("\n")
    except Exception:
        log.exception(
            "Error applying import transforms",
        )
        return dst_fallback


def _schedule_module_imports(
    gatherer: GatherImportsVisitor,
    collector: DottedImportCollector,
    ctx: CodemodContext,
    dst_names: set[str],
) -> None:
    """Schedule module-level imports for addition/removal."""
    for mod in gatherer.module_imports:
        if mod == "__future__":
            continue
        bound_name = mod.split(".")[0]
        if bound_name in dst_names and mod not in collector.imports:
            AddImportsVisitor.add_needed_import(ctx, mod)
        RemoveImportsVisitor.remove_unused_import(ctx, mod)


def _schedule_object_imports(  # noqa: C901, PLR0913
    gatherer: GatherImportsVisitor,
    collector: DottedImportCollector,
    ctx: CodemodContext,
    fqn_set: set[str],
    project_root: Path,
    dst_names: set[str],
) -> None:
    """Schedule from-imports for addition/removal."""
    aliased_objects: set[str] = set()
    for mod, alias_pairs in gatherer.alias_mapping.items():
        for pair in alias_pairs:
            if pair[0] and pair[1]:
                aliased_objects.add(
                    f"{mod}.{pair[0]}",
                )

    for mod, obj_seq in gatherer.object_mapping.items():
        for obj in obj_seq:
            fqn = f"{mod}.{obj}"
            if fqn in fqn_set or ctx.full_module_name == mod:
                continue
            if fqn in aliased_objects:
                continue

            if obj == "*":
                for sym in resolve_star_import(
                    mod,
                    project_root,
                ):
                    sym_fqn = f"{mod}.{sym}"
                    if (
                        sym in dst_names
                        and sym_fqn not in fqn_set
                        and sym_fqn not in collector.imports
                    ):
                        AddImportsVisitor.add_needed_import(
                            ctx,
                            mod,
                            sym,
                        )
                    RemoveImportsVisitor.remove_unused_import(
                        ctx,
                        mod,
                        sym,
                    )
            else:
                if (
                    mod == "__future__" or obj in dst_names
                ) and fqn not in collector.imports:
                    AddImportsVisitor.add_needed_import(
                        ctx,
                        mod,
                        obj,
                    )
                RemoveImportsVisitor.remove_unused_import(
                    ctx,
                    mod,
                    obj,
                )


def _schedule_alias_imports(
    gatherer: GatherImportsVisitor,
    collector: DottedImportCollector,
    ctx: CodemodContext,
    fqn_set: set[str],
    dst_names: set[str],
) -> None:
    """Schedule aliased imports for addition/removal."""
    for mod, asname in gatherer.module_aliases.items():
        if not asname:
            continue
        if asname in dst_names and f"{mod}.{asname}" not in collector.imports:
            AddImportsVisitor.add_needed_import(
                ctx,
                mod,
                asname=asname,
            )
        RemoveImportsVisitor.remove_unused_import(
            ctx,
            mod,
            asname=asname,
        )

    for (
        mod,
        alias_pairs,
    ) in gatherer.alias_mapping.items():
        for pair in alias_pairs:
            if f"{mod}.{pair[0]}" in fqn_set:
                continue
            if not pair[0] or not pair[1]:
                continue
            if (
                pair[1] in dst_names
                and f"{mod}.{pair[1]}" not in collector.imports
            ):
                AddImportsVisitor.add_needed_import(
                    ctx,
                    mod,
                    pair[0],
                    asname=pair[1],
                )
            RemoveImportsVisitor.remove_unused_import(
                ctx,
                mod,
                pair[0],
                asname=pair[1],
            )
