"""Collect, replace, and insert global definitions.

Handles module-level functions, assignments, and
miscellaneous statements that need to be transferred
between source and destination modules.
"""

from __future__ import annotations

from itertools import chain
from typing import TYPE_CHECKING, Union

import libcst as cst

if TYPE_CHECKING:
    from collections.abc import Sequence


def find_insertion_index_after_imports(
    node: cst.Module,
) -> int:
    """Find the position after the last import statement."""
    insert_index = 0
    for i, stmt in enumerate(node.body):
        is_top_level_import = isinstance(
            stmt,
            cst.SimpleStatementLine,
        ) and any(
            isinstance(child, (cst.Import, cst.ImportFrom))
            for child in stmt.body
        )
        is_conditional_import = isinstance(
            stmt,
            cst.If,
        ) and all(
            isinstance(inner, cst.SimpleStatementLine)
            and all(
                isinstance(
                    child,
                    (cst.Import, cst.ImportFrom),
                )
                for child in inner.body
            )
            for inner in stmt.body.body
        )
        if is_top_level_import or is_conditional_import:
            insert_index = i + 1
        if isinstance(stmt, (cst.ClassDef, cst.FunctionDef)):
            break
    return insert_index


def collect_referenced_names(
    node: cst.CSTNode,
) -> set[str]:
    """Collect all names referenced in a CST node."""
    names: set[str] = set()

    def _collect(n: cst.CSTNode) -> None:
        """Recursively collect Name node values."""
        if isinstance(n, cst.Name):
            names.add(n.value)
        for child in n.children:
            _collect(child)

    _collect(node)
    return names


# -- Global function collectors/transformers -----------------------


class GlobalFunctionCollector(cst.CSTVisitor):
    """Collect module-level function definitions."""

    def __init__(self) -> None:
        """Initialize with empty function collection."""
        super().__init__()
        self.functions: dict[str, cst.FunctionDef] = {}
        self.function_order: list[str] = []

    def visit_FunctionDef(  # noqa: N802
        self,
        node: cst.FunctionDef,
    ) -> bool | None:
        """Record function and skip its body."""
        name = node.name.value
        self.functions[name] = node
        if name not in self.function_order:
            self.function_order.append(name)
        return False

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


class GlobalFunctionTransformer(cst.CSTTransformer):
    """Add/replace module-level functions from new code."""

    def __init__(
        self,
        new_functions: dict[str, cst.FunctionDef],
        new_function_order: list[str],
    ) -> None:
        """Initialize with new function definitions."""
        super().__init__()
        self.new_functions = new_functions
        self.new_function_order = new_function_order
        self.processed_functions: set[str] = set()

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

    def leave_FunctionDef(  # noqa: N802
        self,
        original_node: cst.FunctionDef,
        updated_node: cst.FunctionDef,
    ) -> cst.FunctionDef:
        """Replace function if it exists in new code."""
        name = original_node.name.value
        if name in self.new_functions:
            self.processed_functions.add(name)
            return self.new_functions[name]
        return updated_node

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

    def leave_Module(  # noqa: N802
        self,
        original_node: cst.Module,
        updated_node: cst.Module,
    ) -> cst.Module:
        """Append new functions not in the module."""
        new_statements = list(updated_node.body)
        functions_to_append = [
            self.new_functions[name]
            for name in self.new_function_order
            if name not in self.processed_functions
            and name in self.new_functions
        ]
        if functions_to_append:
            insert_index = find_insertion_index_after_imports(
                updated_node,
            )
            for i, stmt in enumerate(new_statements):
                if isinstance(
                    stmt,
                    (cst.FunctionDef, cst.ClassDef),
                ):
                    insert_index = i + 1
            function_nodes = [
                func.with_changes(
                    leading_lines=[
                        cst.EmptyLine(),
                        *func.leading_lines,
                    ],
                )
                for func in functions_to_append
            ]
            new_statements = list(
                chain(
                    new_statements[:insert_index],
                    function_nodes,
                    new_statements[insert_index:],
                ),
            )
        return updated_node.with_changes(
            body=new_statements,
        )


# -- Global assignment collectors/transformers --------------------


class GlobalAssignmentCollector(cst.CSTVisitor):
    """Collect global assignment statements."""

    def __init__(self) -> None:
        """Initialize with empty assignment collection."""
        super().__init__()
        self.assignments: dict[str, cst.Assign | cst.AnnAssign] = {}
        self.assignment_order: list[str] = []
        self.if_else_depth = 0

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

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

    def visit_If(  # noqa: N802
        self,
        node: cst.If,
    ) -> bool | None:
        """Track conditional nesting depth."""
        self.if_else_depth += 1
        return True

    def leave_If(  # noqa: N802
        self,
        original_node: cst.If,
    ) -> None:
        """Track conditional nesting depth."""
        self.if_else_depth -= 1

    def visit_Assign(  # noqa: N802
        self,
        node: cst.Assign,
    ) -> bool | None:
        """Record top-level assignments."""
        if self.if_else_depth == 0:
            for target in node.targets:
                if isinstance(target.target, cst.Name):
                    name = target.target.value
                    self.assignments[name] = node
                    if name not in self.assignment_order:
                        self.assignment_order.append(
                            name,
                        )
        return True

    def visit_AnnAssign(  # noqa: N802
        self,
        node: cst.AnnAssign,
    ) -> bool | None:
        """Record top-level annotated assignments."""
        if (
            self.if_else_depth == 0
            and isinstance(node.target, cst.Name)
            and node.value is not None
        ):
            name = node.target.value
            self.assignments[name] = node
            if name not in self.assignment_order:
                self.assignment_order.append(name)
        return True


def _partition_new_assignments(
    to_append: list[tuple[str, cst.Assign | cst.AnnAssign]],
    module_defined_names: set[str],
) -> tuple[
    list[tuple[str, cst.Assign | cst.AnnAssign]],
    list[tuple[str, cst.Assign | cst.AnnAssign]],
]:
    """Split assignments into import-safe and def-dependent."""
    after_imports: list[tuple[str, cst.Assign | cst.AnnAssign]] = []
    after_defs: list[tuple[str, cst.Assign | cst.AnnAssign]] = []
    for name, assignment in to_append:
        if (
            isinstance(
                assignment,
                (cst.Assign, cst.AnnAssign),
            )
            and assignment.value is not None
        ):
            refs = collect_referenced_names(
                assignment.value,
            )
            if refs & module_defined_names:
                after_defs.append((name, assignment))
            else:
                after_imports.append(
                    (name, assignment),
                )
        else:
            after_imports.append((name, assignment))
    return after_imports, after_defs


_BodyStmt = Union[cst.SimpleStatementLine, cst.BaseCompoundStatement]


def _insert_assignment_lines(
    stmts: Sequence[_BodyStmt],
    assignments: list[tuple[str, cst.Assign | cst.AnnAssign]],
    idx: int,
) -> list[_BodyStmt]:
    """Insert assignment statements at *idx*."""
    lines = [
        cst.SimpleStatementLine(
            [a],
            leading_lines=[cst.EmptyLine()],
        )
        for _, a in assignments
    ]
    return list(
        chain(stmts[:idx], lines, stmts[idx:]),
    )


class GlobalAssignmentTransformer(cst.CSTTransformer):
    """Replace/add global assignments from new code."""

    def __init__(
        self,
        new_assignments: dict[str, cst.Assign | cst.AnnAssign],
        new_assignment_order: list[str],
    ) -> None:
        """Initialize with new assignments."""
        super().__init__()
        self.new_assignments = new_assignments
        self.new_assignment_order = new_assignment_order
        self.processed_assignments: set[str] = set()
        self.if_else_depth = 0

    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:
        """Track conditional nesting depth."""
        self.if_else_depth += 1

    def leave_If(  # noqa: N802
        self,
        original_node: cst.If,
        updated_node: cst.If,
    ) -> cst.If:
        """Track conditional nesting depth."""
        self.if_else_depth -= 1
        return updated_node

    def leave_Assign(  # noqa: N802
        self,
        original_node: cst.Assign,
        updated_node: cst.Assign,
    ) -> (
        cst.BaseSmallStatement
        | cst.FlattenSentinel[cst.BaseSmallStatement]
        | cst.RemovalSentinel
    ):
        """Replace matching assignments."""
        if self.if_else_depth > 0:
            return updated_node
        for target in original_node.targets:
            if isinstance(target.target, cst.Name):
                name = target.target.value
                if name in self.new_assignments:
                    self.processed_assignments.add(
                        name,
                    )
                    return self.new_assignments[name]
        return updated_node

    def leave_AnnAssign(  # noqa: N802
        self,
        original_node: cst.AnnAssign,
        updated_node: cst.AnnAssign,
    ) -> (
        cst.BaseSmallStatement
        | cst.FlattenSentinel[cst.BaseSmallStatement]
        | cst.RemovalSentinel
    ):
        """Replace matching annotated assignments."""
        if self.if_else_depth > 0:
            return updated_node
        if isinstance(original_node.target, cst.Name):
            name = original_node.target.value
            if name in self.new_assignments:
                self.processed_assignments.add(name)
                return self.new_assignments[name]
        return updated_node

    def leave_Module(  # noqa: N802
        self,
        original_node: cst.Module,
        updated_node: cst.Module,
    ) -> cst.Module:
        """Add new assignments not in the module."""
        new_stmts = list(updated_node.body)
        to_append = [
            (name, self.new_assignments[name])
            for name in self.new_assignment_order
            if name not in self.processed_assignments
            and name in self.new_assignments
        ]
        if not to_append:
            return updated_node.with_changes(
                body=new_stmts,
            )

        module_defined_names: set[str] = set()
        for stmt in new_stmts:
            if isinstance(
                stmt,
                (cst.ClassDef, cst.FunctionDef),
            ):
                module_defined_names.add(
                    stmt.name.value,
                )

        after_imports, after_defs = _partition_new_assignments(
            to_append,
            module_defined_names,
        )

        if after_imports:
            idx = find_insertion_index_after_imports(
                updated_node,
            )
            new_stmts = _insert_assignment_lines(
                new_stmts,
                after_imports,
                idx,
            )

        if after_defs:
            idx = find_insertion_index_after_imports(
                cst.Module(body=new_stmts),
            )
            for i, stmt in enumerate(new_stmts):
                if isinstance(
                    stmt,
                    (cst.FunctionDef, cst.ClassDef),
                ):
                    idx = i + 1
            new_stmts = _insert_assignment_lines(
                new_stmts,
                after_defs,
                idx,
            )

        return updated_node.with_changes(
            body=new_stmts,
        )


# -- Global statement collectors/transformers ---------------------


class GlobalStatementCollector(cst.CSTVisitor):
    """Collect module-level non-import, non-assignment stmts."""

    def __init__(self) -> None:
        """Initialize with empty statement list."""
        super().__init__()
        self.global_statements: list[cst.SimpleStatementLine] = []

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

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

    def visit_SimpleStatementLine(  # noqa: N802
        self,
        node: cst.SimpleStatementLine,
    ) -> None:
        """Record non-import, non-assignment statements."""
        for statement in node.body:
            if not isinstance(
                statement,
                (
                    cst.Import,
                    cst.ImportFrom,
                    cst.Assign,
                    cst.AnnAssign,
                ),
            ):
                self.global_statements.append(node)
                break


class GlobalStatementTransformer(cst.CSTTransformer):
    """Append global statements at end of module."""

    def __init__(
        self,
        global_statements: list[cst.SimpleStatementLine],
    ) -> None:
        """Initialize with statements to append."""
        super().__init__()
        self.global_statements = global_statements

    def leave_Module(  # noqa: N802
        self,
        original_node: cst.Module,
        updated_node: cst.Module,
    ) -> cst.Module:
        """Append statements after all definitions."""
        if not self.global_statements:
            return updated_node
        new_statements = list(updated_node.body)
        statement_lines = [
            stmt.with_changes(
                leading_lines=[
                    cst.EmptyLine(),
                    *stmt.leading_lines,
                ],
            )
            for stmt in self.global_statements
        ]
        new_statements.extend(statement_lines)
        return updated_node.with_changes(
            body=new_statements,
        )


def extract_global_statements(
    source_code: str,
) -> tuple[cst.Module, list[cst.SimpleStatementLine]]:
    """Extract global statements from source code."""
    module = cst.parse_module(source_code)
    collector = GlobalStatementCollector()
    module.visit(collector)
    return module, collector.global_statements


def add_global_assignments(
    src_module_code: str,
    dst_module_code: str,
) -> str:
    """Add global assignments and functions from *src* to *dst*."""
    src_module, new_global_stmts = extract_global_statements(src_module_code)
    dst_module, existing_global_stmts = extract_global_statements(
        dst_module_code
    )

    unique_global_stmts = []
    for stmt in new_global_stmts:
        if any(
            stmt is existing or stmt.deep_equals(existing)
            for existing in existing_global_stmts
        ):
            continue
        unique_global_stmts.append(stmt)

    assign_collector = GlobalAssignmentCollector()
    src_module.visit(assign_collector)

    src_fn_collector = GlobalFunctionCollector()
    src_module.visit(src_fn_collector)

    dst_fn_collector = GlobalFunctionCollector()
    dst_module.visit(dst_fn_collector)

    new_functions = {
        name: func
        for name, func in (src_fn_collector.functions.items())
        if name not in dst_fn_collector.functions
    }
    new_fn_order = [
        name
        for name in src_fn_collector.function_order
        if name in new_functions
    ]

    if (
        not assign_collector.assignments
        and not new_functions
        and not unique_global_stmts
    ):
        return dst_module_code

    if new_functions:
        dst_module = dst_module.visit(
            GlobalFunctionTransformer(
                new_functions,
                new_fn_order,
            ),
        )

    if assign_collector.assignments:
        dst_module = dst_module.visit(
            GlobalAssignmentTransformer(
                assign_collector.assignments,
                assign_collector.assignment_order,
            ),
        )

    if unique_global_stmts:
        dst_module = dst_module.visit(
            GlobalStatementTransformer(
                unique_global_stmts,
            ),
        )

    return dst_module.code
