"""Import gathering and addition for context extraction."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

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

log = logging.getLogger(__name__)

_SENTINEL = object()


def _needs_future_alias_strip(code: str | cst.Module) -> bool:
    """Fast check for aliased ``__future__`` imports.

    Returns *True* only when the source appears to contain
    ``from __future__ import X as Y``, which is rare in
    practice.  When *False*, the
    :class:`FutureAliasedImportTransformer` can be skipped.
    """
    if isinstance(code, cst.Module):
        code = code.code
    return "__future__" in code and " as " in code


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

    def leave_ImportFrom(  # noqa: N802
        self,
        original_node: cst.ImportFrom,
        updated_node: cst.ImportFrom,
    ) -> (
        cst.BaseSmallStatement
        | cst.FlattenSentinel[cst.BaseSmallStatement]
        | cst.RemovalSentinel
    ):
        """Strip aliased __future__ imports, keeping non-aliased ones."""
        mod = updated_node.module
        if (
            isinstance(mod, cst.Name)
            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
            )
        ):
            kept = [n for n in updated_node.names if n.asname is None]
            if kept:
                return updated_node.with_changes(
                    names=kept,
                )
            return cst.RemoveFromParent()
        return updated_node


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

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

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

    def get_full_dotted_name(
        self,
        expr: cst.BaseExpression,
    ) -> str:
        """Convert a CST expression to a dotted name."""
        if isinstance(expr, cst.Name):
            return expr.value
        if isinstance(expr, cst.Attribute):
            prefix = self.get_full_dotted_name(expr.value)
            return f"{prefix}.{expr.attr.value}"
        return ""

    def _handle_import(
        self,
        node: cst.Import,
    ) -> None:
        """Record dotted names from a bare import statement."""
        if isinstance(node.names, cst.ImportStar):
            return
        for alias in node.names:
            mod = self.get_full_dotted_name(alias.name)
            if alias.asname is not None:
                name_node = alias.asname.name
                if isinstance(name_node, cst.Name):
                    asname: str = name_node.value
                    self.imports.add(
                        mod if mod == asname else f"{mod}.{asname}"
                    )
                else:
                    self.imports.add(mod)
            else:
                self.imports.add(mod)

    def _handle_import_from(
        self,
        node: cst.ImportFrom,
    ) -> None:
        """Record dotted names from a from-import statement."""
        if node.module is None:
            return
        mod = 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 is not None:
                asname_node = alias.asname.name
                if isinstance(asname_node, cst.Name):
                    asname_str = asname_node.value
                else:
                    asname_str = name
            else:
                asname_str = name
            self.imports.add(f"{mod}.{asname_str}")

    def _collect_imports_from_block(
        self,
        block: cst.IndentedBlock | cst.Module,
    ) -> None:
        """Walk a block's simple statements for import nodes."""
        for stmt in block.body:
            if not isinstance(
                stmt,
                cst.SimpleStatementLine,
            ):
                continue
            for child in stmt.body:
                if isinstance(child, cst.Import):
                    self._handle_import(child)
                elif isinstance(child, cst.ImportFrom):
                    self._handle_import_from(child)

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

    def visit_FunctionDef(  # noqa: N802
        self,
        node: cst.FunctionDef,
    ) -> bool:
        """Do not descend into functions."""
        return False

    def visit_ClassDef(  # noqa: N802
        self,
        node: cst.ClassDef,
    ) -> bool:
        """Do not descend into classes."""
        return False

    def visit_If(  # noqa: N802
        self,
        node: cst.If,
    ) -> None:
        """Collect imports from conditional 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 from try blocks."""
        if isinstance(node.body, cst.IndentedBlock):
            self._collect_imports_from_block(node.body)


def gather_source_imports(
    src_module_code: str | cst.Module,
    src_path: Path,
    project_root: Path,
) -> GatherImportsVisitor | None:
    """Gather imports from a source module.

    Returns *None* when the source has no module-level
    imports.  Call once per source file and pass the result
    to *add_needed_imports_from_module* via
    *gathered_imports*.
    """
    mod_and_pkg = calculate_module_and_package(
        project_root,
        src_path,
    )
    try:
        if isinstance(src_module_code, cst.Module):
            src_module = src_module_code
        else:
            src_module = cst.parse_module(src_module_code)

        if _needs_future_alias_strip(src_module_code):
            src_module = src_module.visit(
                FutureAliasedImportTransformer(),
            )

        has_imports = any(
            isinstance(s, (cst.Import, cst.ImportFrom))
            for stmt in src_module.body
            if isinstance(stmt, cst.SimpleStatementLine)
            for s in stmt.body
        )
        if not has_imports:
            return None

        gatherer = GatherImportsVisitor(
            CodemodContext(
                filename=src_path.name,
                full_module_name=mod_and_pkg.name,
                full_package_name=mod_and_pkg.package,
            ),
        )

        module_level = src_module.with_changes(
            body=[
                s
                for s in src_module.body
                if not isinstance(
                    s,
                    (cst.FunctionDef, cst.ClassDef),
                )
            ],
        )
        module_level.visit(gatherer)

        if (
            not gatherer.module_imports
            and not gatherer.object_mapping
            and not gatherer.module_aliases
            and not gatherer.alias_mapping
        ):
            return None

        return gatherer  # noqa: TRY300
    except Exception:  # noqa: BLE001
        log.debug(
            "Error parsing source module",
            exc_info=True,
        )
        return None


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

    Imports for functions in *helper_fqns* are skipped
    (they are already present in the context).
    """
    if helper_fqns is None:
        helper_fqns = set()

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

    dst_mod_pkg = calculate_module_and_package(
        project_root,
        dst_path,
    )
    dst_ctx = CodemodContext(
        filename=src_path.name,
        full_module_name=dst_mod_pkg.name,
        full_package_name=dst_mod_pkg.package,
    )

    if gathered_imports is _SENTINEL:
        gatherer: GatherImportsVisitor | None = gather_source_imports(
            src_module_code,
            src_path,
            project_root,
        )
    else:
        gatherer = gathered_imports

    if gatherer is None:
        return fallback

    dotted = DottedImportCollector()
    if isinstance(dst_module_code, str):
        try:
            parsed_dst = cst.parse_module(dst_module_code)
        except cst.ParserSyntaxError:
            log.debug(
                "Syntax error in dst module",
                exc_info=True,
            )
            return fallback
    else:
        parsed_dst = dst_module_code

    parsed_dst.visit(dotted)

    try:
        for mod in gatherer.module_imports:
            if mod == "__future__":
                continue
            if mod not in dotted.imports:
                AddImportsVisitor.add_needed_import(
                    dst_ctx,
                    mod,
                )
            RemoveImportsVisitor.remove_unused_import(
                dst_ctx,
                mod,
            )

        aliased_objects: set[str] = set()
        for mod, pairs in gatherer.alias_mapping.items():
            for pair in pairs:
                if pair[0] and pair[1]:
                    aliased_objects.add(
                        f"{mod}.{pair[0]}",
                    )

        for mod, objs in gatherer.object_mapping.items():
            for obj in objs:
                fqn = f"{mod}.{obj}"
                if fqn in helper_fqns or dst_ctx.full_module_name == mod:
                    continue
                if fqn in aliased_objects:
                    continue
                if obj == "*":
                    continue
                if fqn not in dotted.imports:
                    AddImportsVisitor.add_needed_import(
                        dst_ctx,
                        mod,
                        obj,
                    )
                RemoveImportsVisitor.remove_unused_import(
                    dst_ctx,
                    mod,
                    obj,
                )
    except Exception:  # noqa: BLE001
        log.debug(
            "Error adding imports",
            exc_info=True,
        )
        return fallback

    for mod, asname in gatherer.module_aliases.items():
        if not asname:
            continue
        if f"{mod}.{asname}" not in dotted.imports:
            AddImportsVisitor.add_needed_import(
                dst_ctx,
                mod,
                asname=asname,
            )
        RemoveImportsVisitor.remove_unused_import(
            dst_ctx,
            mod,
            asname=asname,
        )

    for mod, pairs in gatherer.alias_mapping.items():
        for pair in pairs:
            if f"{mod}.{pair[0]}" in helper_fqns:
                continue
            if not pair[0] or not pair[1]:
                continue
            if f"{mod}.{pair[1]}" not in dotted.imports:
                AddImportsVisitor.add_needed_import(
                    dst_ctx,
                    mod,
                    pair[0],
                    asname=pair[1],
                )
            RemoveImportsVisitor.remove_unused_import(
                dst_ctx,
                mod,
                pair[0],
                asname=pair[1],
            )

    try:
        result = parsed_dst
        if dst_ctx.scratch.get("AddImportsVisitor"):
            result = AddImportsVisitor(
                dst_ctx,
            ).transform_module(result)
        if dst_ctx.scratch.get("RemoveImportsVisitor"):
            result = RemoveImportsVisitor(
                dst_ctx,
            ).transform_module(result)
        return result.code.lstrip("\n")
    except Exception:  # noqa: BLE001
        log.debug(
            "Error transforming dst module",
            exc_info=True,
        )
        return fallback
