#!/usr/bin/env python3
# --------------------( LICENSE                            )--------------------
# Copyright (c) 2014-2025 Beartype authors.
# See "LICENSE" for further details.

'''
Beartype **module abstract syntax tree (AST) transformers** (i.e., low-level
mixins instrumenting globally scoped import statements prefixing the bodies of
well-typed third-party modules with runtime type-checking dynamically generated
by the :func:`beartype.beartype` decorator).

This private submodule is *not* intended for importation by downstream callers.
'''

# ....................{ IMPORTS                            }....................
from ast import (
    AST,
    Constant,
    Expr,
    ImportFrom,
    Module,
)
from beartype._util.ast.utilastmake import make_node_importfrom

# ....................{ SUBCLASSES                         }....................
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# CAUTION: To improve forward compatibility with the superclass API over which
# we have *NO* control, avoid accidental conflicts by suffixing *ALL* private
# and public attributes of this subclass by "_beartype".
#!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

#FIXME: Unit test us up, please.
class BeartypeNodeTransformerModuleMixin(object):
    '''
    Beartype **module abstract syntax tree (AST) transformers** (i.e., low-level
    mixin instrumenting globally scoped import statements prefixing the bodies
    of well-typed third-party modules with runtime type-checking dynamically
    generated by the :func:`beartype.beartype` decorator).

    '''

    # ..................{ VISITORS ~ module                  }..................
    def visit_Module(self, node: Module) -> Module:
        '''
        Add a new abstract syntax tree (AST) child node to the passed
        **module node** (i.e., node encapsulating the module currently being
        loaded by the
        :class:`beartype.claw._importlib._clawimpload.BeartypeSourceFileLoader`)
        importing various attributes required by lower-level child nodes added
        by subsequent visitor methods defined by this transformer.

        Parameters
        ----------
        node : Module
            Module node to be transformed.

        Returns
        -------
        Module
            That same module node.
        '''

        # 0-based index of an early child node of this parent module node
        # immediately *BEFORE* which to insert one or more statements importing
        # beartype-specific attributes, defaulting to the first child node of
        # this parent module. Specifically, if this module begins with:
        # * Neither a module docstring *NOR* any "from __future__" imports, this
        #   index is guaranteed to be 0.
        # * Only a module docstring but *NO* "from __future__" imports, this
        #   index is guaranteed to be 1.
        # * A module docstring and one or more "from __future__" imports, this
        #   index is guaranteed to be one more than the number of such imports.
        node_index_import_beartype_attrs = 0

        # 0-based index of the last child node of this parent module node.
        #
        # Note that the "node.body" instance variable for module nodes is a list
        # of *ALL* child nodes of this parent module node.
        node_index_last = len(node.body)

        # Child node of this parent module node immediately preceding the output
        # import child node to be added below, defaulting to this parent module
        # node to ensure that the copy_node_metadata() function below *ALWAYS*
        # copies from a valid node (for simplicity).
        node_prev: AST = node

        #FIXME: [SPEED] Refactor into a "while" loop, please. We even localize
        #"node_index_last" above, which is the obvious stop condition. *sigh*
        # For the 0-based index and value of each direct child node of this
        # parent module node...
        #
        # This iteration efficiently finds "node_index_import_beartype_attrs"
        # (i.e., the 0-based index of the first safe position in the list of all
        # child nodes of this parent module node to insert an import statement
        # importing our beartype decorator). Despite superficially appearing to
        # perform a linear search of all "n" child nodes of this module parent
        # node and thus exhibit worst-case O(n) linear time complexity, this
        # iteration is guaranteed to exhibit worst-case O(k) quasi-constant time
        # complexity for "k" the number of "from __future__" import statements
        # prefixing the body of this module. \o/
        for node_prev in node.body:
            # print(f'node_index_import_beartype_attrs [IN]: {node_index_import_beartype_attrs}')

            # If it is *NOT* the case that this child node signifies either...
            if not (
                # A module docstring...
                #
                # If that module defines a docstring, that docstring *MUST* be
                # the first expression of that module. That docstring *MUST* be
                # explicitly found and iterated past to ensure that the import
                # statement added below appears *AFTER* rather than *BEFORE* any
                # docstring. (The latter would destroy the semantics of that
                # docstring by reducing that docstring to an ignorable string.)
                (
                    isinstance(node_prev, Expr) and
                    isinstance(node_prev.value, Constant)
                ) or
                # A future import (i.e., import of the form "from __future__
                # ...") *OR*...
                #
                # If that module performs one or more future imports, these
                # imports *MUST* necessarily be the first non-docstring
                # statement of that module and thus appear *BEFORE* all import
                # statements that are actually imports -- including the import
                # statement added below.
                (
                    isinstance(node_prev, ImportFrom) and
                    node_prev.module == '__future__'
                )
            # Then immediately halt iteration, guaranteeing O(k) runtime.
            ):
                break
            # Else, this child node signifies either a module docstring of
            # future import. In this case, implicitly skip past this child node
            # to the next child node.

            # Insert beartype-specific attributes immediately *AFTER* this node.
            node_index_import_beartype_attrs += 1
        # "node_index_import_beartype_attrs" is now the index of the first safe
        # position in this list to insert output child import nodes below.
        # print(f'node_index_import_beartype_attrs [AFTER]: {node_index_import_beartype_attrs}')
        # print(f'len(node.body): {len(node.body)}')

        #FIXME: Pretty sure this "if" comparison can just be trivially reduced
        #to a loop "else": e.g.,
        #else:  # <-- seriously, that's it.

        # If the 0-based index of an early child node of this parent module node
        # immediately *BEFORE* which to insert one or more statements importing
        # beartype-specific attributes is *NOT* that of the last child node of
        # this parent module node, this module contains one or more semantically
        # meaningful child nodes and is thus non-empty. In this case...
        if node_index_import_beartype_attrs != node_index_last:
            # print('Injecting beartype imports...')

            # Module-scoped import nodes (i.e., child nodes to be inserted under
            # the parent node encapsulating the currently visited submodule in
            # the AST for that module).
            #
            # Note that:
            # * The original attributes are imported into the currently visited
            #   submodule under obfuscated beartype-specific names,
            #   significantly reducing the likelihood of a namespace collision
            #   with existing attributes of the same name in that submodule.
            # * These nodes are intentionally *NOT* generalized into global
            #   constants. In theory, doing so would reduce space and time
            #   complexity by enabling efficient reuse here. In practice, doing
            #   so would also be fundamentally wrong; these nodes are
            #   subsequently modified to respect the source code metadata (e.g.,
            #   line numbers) of this AST module parent node, which prevents
            #   such trivial reuse. Although we could further attempt to
            #   circumvent that by shallowly or deeply copying from global
            #   constants, both the copy() and deepcopy() functions defined by
            #   the standard "copy" module are pure-Python and thus shockingly
            #   slow -- which defeats the purpose.

            # Node importing all beartype-specific attributes explicitly
            # imported and implicitly exported by our private
            # "beartype.claw._ast.clawaststar" submodule, comprising the set of
            # all attributes required by code dynamically injected into this AST
            # by this AST transformer.
            node_import_all = make_node_importfrom(
                module_name='beartype.claw._ast._clawaststar',
                source_attr_name='*',
                node_sibling=node_prev,
            )

            #FIXME: Explain this slice operation. Super-bizarre, honestly.
            # Insert these output child import nodes at this safe position of
            # the list of all child nodes of this parent module node.
            #
            # Note that this syntax efficiently (albeit unreadably) inserts
            # these output child import nodes at the desired index (in this
            # arbitrary order) of this parent module node.
            node.body[node_index_import_beartype_attrs:0] = (node_import_all,)
        # Else, the 0-based index of an early child node of this parent module
        # node immediately *BEFORE* which to insert one or more statements
        # importing beartype-specific attributes is that of the last child node
        # of this parent module node. In this case, this module contains *NO*
        # semantically meaningful child nodes and is thus effectively empty.
        # In this case, silently reduce to a noop. This edge case is *EXTREMELY*
        # uncommon and thus *NOT* optimized for (either here or elsewhere).
        #
        # Note that this edge cases cleanly matches:
        # * Syntactically empty modules containing only zero or more whitespace
        #   characters and zero or more inline comments.
        # * Syntactically non-empty modules containing only a prefacing module
        #   docstring and/or one or more "from __future__" import statements.
        #   Semantically, these sorts of modules are effectively empty as well.

        # Recursively transform *ALL* child nodes of this parent module node.
        node = self.generic_visit(node)  # type: ignore[attr-defined]

        # #FIXME: Conditionally perform this logic if "conf.is_debug", please.
        # print(
        #     f'Module abstract syntax tree (AST) transformed by @beartype to:\n\n'
        #     f'{get_node_repr_indented(node)}'
        # )

        # Return this transformed module node.
        return node
