# -*- coding: utf-8; -*-
"""Fix `ctx` attributes and source location info in an AST."""

__all__ = ["fix_ctx", "fix_locations"]

from ast import (Load, Store, Del,
                 Assign, AnnAssign, AugAssign,
                 Attribute, Subscript,
                 comprehension,
                 For, AsyncFor,
                 withitem,
                 Delete,
                 iter_child_nodes)
from copy import copy

from . import walkers

try:  # Python 3.8+
    from ast import NamedExpr
except ImportError:
    class _NoSuchNodeType:
        pass
    NamedExpr = _NoSuchNodeType


class _CtxFixer(walkers.ASTTransformer):
    def __init__(self, *, copy_seen_nodes):
        super().__init__(ctxclass=Load)
        self.copy_seen_nodes = copy_seen_nodes

    def reset(self, **bindings):
        super().reset(**bindings)
        self._seen = set()

    def transform(self, tree):
        tree = self._fix_one(tree)
        self._setup_subtree_contexts(tree)
        return self.generic_visit(tree)

    def _fix_one(self, tree):
        """Fix one `ctx` attribute, using the currently active ctx class."""
        if "ctx" in type(tree)._fields:
            if self.copy_seen_nodes:
                # Shallow-copy `tree` if already seen. This mode is used in the
                # global postprocess pass. Note only nodes whose `ctx` we have
                # already updated are considered seen.
                if id(tree) in self._seen:
                    tree = copy(tree)
            tree.ctx = self.state.ctxclass()
            self._seen.add(id(tree))
        return tree

    def _setup_subtree_contexts(self, tree):
        """Autoselect correct `ctx` class for subtrees of `tree`."""
        # The default ctx class is `Load`. We set up any `Store` and `Del`, as
        # well as any `Load` for trees that may appear inside others that are
        # set up as `Store` or `Del` (that mainly concerns expressions).
        tt = type(tree)
        if tt is Assign:
            self.withstate(tree.targets, ctxclass=Store)
            self.withstate(tree.value, ctxclass=Load)
        elif tt is AnnAssign:
            self.withstate(tree.target, ctxclass=Store)
            self.withstate(tree.annotation, ctxclass=Load)
            if tree.value:
                self.withstate(tree.value, ctxclass=Load)
        elif tt is NamedExpr:
            self.withstate(tree.target, ctxclass=Store)
            self.withstate(tree.value, ctxclass=Load)
        elif tt is AugAssign:
            # `AugStore` and `AugLoad` are for internal use only, not even
            # meant to be exposed to the user; the compiler expects `Store`
            # and `Load` here. https://bugs.python.org/issue39988
            self.withstate(tree.target, ctxclass=Store)
            self.withstate(tree.value, ctxclass=Load)

        elif tt is Attribute:
            # The tree's own `ctx` can be whatever, but `value` always has `Load`.
            self.withstate(tree.value, ctxclass=Load)
        elif tt is Subscript:
            # The tree's own `ctx` can be whatever, but `value` and `slice` always have `Load`.
            self.withstate(tree.value, ctxclass=Load)
            self.withstate(tree.slice, ctxclass=Load)

        elif tt is comprehension:
            self.withstate(tree.target, ctxclass=Store)
            self.withstate(tree.iter, ctxclass=Load)
            self.withstate(tree.ifs, ctxclass=Load)

        elif tt in (For, AsyncFor):
            self.withstate(tree.target, ctxclass=Store)
            self.withstate(tree.iter, ctxclass=Load)
        elif tt is withitem:
            self.withstate(tree.context_expr, ctxclass=Load)
            self.withstate(tree.optional_vars, ctxclass=Store)

        elif tt is Delete:
            self.withstate(tree.targets, ctxclass=Del)


def fix_ctx(tree, *, copy_seen_nodes):
    """Fix `ctx` attributes in `tree`.

    If `copy_seen_nodes=True`, then, if the same AST node instance appears
    multiple times, and the node requires a `ctx`, shallow-copy it the second
    and further times it is encountered. This prevents problems if the same
    node instance has been spliced into two or more positions that require
    different `ctx`.

    Modifies `tree` in-place. For convenience, returns the modified `tree`.
    """
    return _CtxFixer(copy_seen_nodes=copy_seen_nodes).visit(tree)


def fix_locations(tree, reference_node, *, mode):
    """Like `ast.fix_missing_locations`, but customized for a macro expander.

    Differences:

      - If `reference_node` has source no location info, return immediately (no-op).
      - If `tree is None`, return immediately (no-op).
      - If `tree` is a `list` of AST nodes, loop over it.

    The `mode` parameter:

      - If `mode="reference"`, populate any missing location info by
        copying it from `reference_node`. Always use the same reference info.

        Good when expanding a macro invocation, to set the source location
        of any macro-generated nodes to that of the macro invocation node.

      - If `mode="update"`, behave exactly like `ast.fix_missing_locations`,
        except that at the top level of `tree`, initialize `lineno` and
        `col_offset` from `reference_node` (instead of using `1` and `0`
        like `ast.fix_missing_locations` does).

        So if a node is missing location info, copy the current reference info
        in, but if it has location info, then update the reference info.

        Good for general use.

      - If `mode="overwrite"`, copy location info from `reference_node`,
        regardless of if the target node already has it.

        Good when `tree` is a code template that comes from another file,
        so that any line numbers already in the AST would be misleading
        at the use site (because they point to lines in that other file).

    Modifies `tree` in-place. For convenience, returns the modified `tree`.
    """
    if not (hasattr(reference_node, "lineno") and hasattr(reference_node, "col_offset")):
        return tree
    def _fix(tree, lineno, col_offset):
        if tree is None:
            return
        if isinstance(tree, list):
            for elt in tree:
                _fix(elt, lineno, col_offset)
            return
        if "lineno" in tree._attributes:
            if mode == "overwrite":
                tree.lineno = lineno
            else:
                if not hasattr(tree, "lineno"):
                    tree.lineno = lineno
                elif mode == "update":
                    lineno = tree.lineno
        if "col_offset" in tree._attributes:
            if mode == "overwrite":
                tree.col_offset = col_offset
            else:
                if not hasattr(tree, "col_offset"):
                    tree.col_offset = col_offset
                elif mode == "update":
                    col_offset = tree.col_offset
        for child in iter_child_nodes(tree):
            _fix(child, lineno, col_offset)
    _fix(tree, reference_node.lineno, reference_node.col_offset)
    return tree
