import sys as _sys
import ast as _ast
from ast import boolop, cmpop, excepthandler, expr, expr_context, operator
from ast import slice, stmt, unaryop, mod, AST
from ast import iter_child_nodes, walk

try:
    from ast import TypeIgnore
except ImportError:
    class TypeIgnore(AST):
        pass


def _make_node(Name, Fields, Attributes, Bases):
    def create_node(self, *args, **kwargs):
        nbparam = len(args) + len(kwargs)
        assert nbparam in (0, len(Fields)), \
            "Bad argument number for {}: {}, expecting {}".\
            format(Name, nbparam, len(Fields))
        self._fields = Fields
        self._attributes = Attributes
        for argname, argval in zip(self._fields, args):
            setattr(self, argname, argval)
        for argname, argval in kwargs.items():
            assert argname in Fields, \
                    "Invalid Keyword argument for {}: {}".format(Name, argname)
            setattr(self, argname, argval)

    setattr(_sys.modules[__name__],
            Name,
            type(Name,
                 Bases,
                 {'__init__': create_node}))


_nodes = (
    # mod
    ('Module', (('body', 'type_ignores'), (), (mod,))),
    ('Interactive', (('body',), (), (mod,))),
    ('Expression', (('body',), (), (mod,))),
    ('FunctionType', (('argtypes', 'returns'), (), (mod,))),
    ('Suite', (('body',), (), (mod,))),

    # stmt
    ('FunctionDef', (('name', 'args', 'body', 'decorator_list', 'returns',
                      'type_comment'),
                     ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                     (stmt,))),
    ('AsyncFunctionDef', (('name', 'args', 'body',
                           'decorator_list', 'returns',
                           'type_comment'),
                          ('lineno', 'col_offset',
                           'end_lineno', 'end_col_offset',),
                          (stmt,))),
    ('ClassDef', (('name', 'bases', 'keywords', 'body', 'decorator_list',),
                  ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                  (stmt,))),
    ('Return', (('value',),
                ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                (stmt,))),
    ('Delete', (('targets',),
                ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                (stmt,))),
    ('Assign', (('targets', 'value',),
                ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                (stmt,))),
    ('AugAssign', (('target', 'op', 'value',),
                   ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                   (stmt,))),
    ('AnnAssign', (('target', 'annotation', 'value', 'simple',),
                   ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                   (stmt,))),
    ('Print', (('dest', 'values', 'nl',),
               ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (stmt,))),
    ('For', (('target', 'iter', 'body', 'orelse', 'type_comment'),
             ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
             (stmt,))),
    ('AsyncFor', (('target', 'iter', 'body', 'orelse', 'type_comment'),
                  ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                  (stmt,))),
    ('While', (('test', 'body', 'orelse',),
               ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (stmt,))),
    ('If', (('test', 'body', 'orelse',),
            ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
            (stmt,))),
    ('With', (('items', 'body', 'type_comment'),
              ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
              (stmt,))),
    ('AsyncWith', (('items', 'body', 'type_comment'),
                   ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                   (stmt,))),
    ('Raise', (('exc', 'cause',),
               ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (stmt,))),
    ('Try', (('body', 'handlers', 'orelse', 'finalbody',),
             ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
             (stmt,))),
    ('Assert', (('test', 'msg',),
                ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                (stmt,))),
    ('Import', (('names',),
                ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                (stmt,))),
    ('ImportFrom', (('module', 'names', 'level',),
                    ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                    (stmt,))),
    ('Exec', (('body', 'globals', 'locals',),
              ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
              (stmt,))),
    ('Global', (('names',),
                ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                (stmt,))),
    ('Nonlocal', (('names',),
                  ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                  (stmt,))),
    ('Expr', (('value',),
              ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
              (stmt,))),
    ('Pass', ((), ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
              (stmt,))),
    ('Break', ((), ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (stmt,))),
    ('Continue', ((),
                  ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                  (stmt,))),

    # expr

    ('BoolOp', (('op', 'values',),
                ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                (expr,))),
    ('BinOp', (('left', 'op', 'right',),
               ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (expr,))),
    ('UnaryOp', (('op', 'operand',),
                 ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                 (expr,))),
    ('Lambda', (('args', 'body',),
                ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                (expr,))),
    ('IfExp', (('test', 'body', 'orelse',),
               ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (expr,))),
    ('Dict', (('keys', 'values',),
              ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
              (expr,))),
    ('Set', (('elts',),
             ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
             (expr,))),
    ('ListComp', (('elt', 'generators',),
                  ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                  (expr,))),
    ('SetComp', (('elt', 'generators',),
                 ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                 (expr,))),
    ('DictComp', (('key', 'value', 'generators',),
                  ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                  (expr,))),
    ('GeneratorExp', (('elt', 'generators',),
                      ('lineno', 'col_offset',
                       'end_lineno', 'end_col_offset',),
                      (expr,))),
    ('Await', (('value',),
               ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (expr,))),
    ('Yield', (('value',),
               ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (expr,))),
    ('YieldFrom', (('value',),
                   ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                   (expr,))),
    ('Compare', (('left', 'ops', 'comparators',),
                 ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                 (expr,))),
    ('Call', (('func', 'args', 'keywords',),
              ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
              (expr,))),
    ('Repr', (('value',),
              ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
              (expr,))),
    ('FormattedValue', (('value', 'conversion', 'format_spec',),
                        ('lineno', 'col_offset',
                         'end_lineno', 'end_col_offset',),
                        (expr,))),
    ('JoinedStr', (('values',),
                   ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                   (expr,))),
    ('Constant', (('value', 'kind'),
                  ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                  (expr,))),
    ('Attribute', (('value', 'attr', 'ctx',),
                   ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                   (expr,))),
    ('Subscript', (('value', 'slice', 'ctx',),
                   ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                   (expr,))),
    ('Starred', (('value', 'ctx',),
                 ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
                 (expr,))),
    ('Name', (('id', 'ctx', 'annotation', 'type_comment'),
              ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
              (expr,))),
    ('List', (('elts', 'ctx',),
              ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
              (expr,))),
    ('Tuple', (('elts', 'ctx',),
               ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (expr,))),

    # expr_context
    ('Load', ((), (), (expr_context,))),
    ('Store', ((), (), (expr_context,))),
    ('Del', ((), (), (expr_context,))),
    ('AugLoad', ((), (), (expr_context,))),
    ('AugStore', ((), (), (expr_context,))),
    ('Param', ((), (), (expr_context,))),

    # slice
    ('Slice', (('lower', 'upper', 'step'),
               ('lineno', 'col_offset', 'end_lineno', 'end_col_offset',),
               (slice,))),

    # boolop
    ('And', ((), (), (boolop,))),
    ('Or', ((), (), (boolop,))),

    # operator
    ('Add', ((), (), (operator,))),
    ('Sub', ((), (), (operator,))),
    ('Mult', ((), (), (operator,))),
    ('MatMult', ((), (), (operator,))),
    ('Div', ((), (), (operator,))),
    ('Mod', ((), (), (operator,))),
    ('Pow', ((), (), (operator,))),
    ('LShift', ((), (), (operator,))),
    ('RShift', ((), (), (operator,))),
    ('BitOr', ((), (), (operator,))),
    ('BitXor', ((), (), (operator,))),
    ('BitAnd', ((), (), (operator,))),
    ('FloorDiv', ((), (), (operator,))),

    # unaryop
    ('Invert', ((), (), (unaryop, AST,))),
    ('Not', ((), (), (unaryop, AST,))),
    ('UAdd', ((), (), (unaryop, AST,))),
    ('USub', ((), (), (unaryop, AST,))),

    # cmpop
    ('Eq', ((), (), (cmpop,))),
    ('NotEq', ((), (), (cmpop,))),
    ('Lt', ((), (), (cmpop,))),
    ('LtE', ((), (), (cmpop,))),
    ('Gt', ((), (), (cmpop,))),
    ('GtE', ((), (), (cmpop,))),
    ('Is', ((), (), (cmpop,))),
    ('IsNot', ((), (), (cmpop,))),
    ('In', ((), (), (cmpop,))),
    ('NotIn', ((), (), (cmpop,))),

    # comprehension
    ('comprehension', (('target', 'iter', 'ifs', 'is_async'), (), (AST,))),

    # excepthandler
    ('ExceptHandler', (('type', 'name', 'body'),
                       ('lineno', 'col_offset',
                        'end_lineno', 'end_col_offset'),
                       (excepthandler,))),

    # arguments
    ('arguments', (('args', 'posonlyargs', 'vararg', 'kwonlyargs',
                    'kw_defaults', 'kwarg', 'defaults'), (), (AST,))),

    # keyword
    ('keyword', (('arg', 'value'),
                 ('lineno', 'col_offset', 'end_lineno', 'end_col_offset'),
                 (AST,))),

    # alias
    ('alias', (('name', 'asname'), (), (AST,))),

    # withitem
    ('withitem', (('context_expr', 'optional_vars'), (), (AST,))),

    # type_ignore
    ('type_ignore', ((), ('lineno', 'tag'), (TypeIgnore,))),
    )

for name, descr in _nodes:
    _make_node(name, *descr)

if _sys.version_info.major == 2:
    from .ast2 import ast_to_gast, gast_to_ast
if _sys.version_info.major == 3:
    from .ast3 import ast_to_gast, gast_to_ast


def parse(*args, **kwargs):
    return ast_to_gast(_ast.parse(*args, **kwargs))


def literal_eval(node_or_string):
    if isinstance(node_or_string, AST):
        node_or_string = gast_to_ast(node_or_string)
    return _ast.literal_eval(node_or_string)


def get_docstring(node, clean=True):
    if not isinstance(node, (FunctionDef, ClassDef, Module)):
        raise TypeError("%r can't have docstrings" % node.__class__.__name__)
    if node.body and isinstance(node.body[0], Expr) and \
       isinstance(node.body[0].value, Constant):
        if clean:
            import inspect
            holder = node.body[0].value
            return inspect.cleandoc(getattr(holder, holder._fields[0]))
        return node.body[0].value.s


# the following are directly imported from python3.8's Lib/ast.py  #

def copy_location(new_node, old_node):
    """
    Copy source location (`lineno`, `col_offset`, `end_lineno`, and
    `end_col_offset` attributes) from *old_node* to *new_node* if possible,
    and return *new_node*.
    """
    for attr in 'lineno', 'col_offset', 'end_lineno', 'end_col_offset':
        if attr in old_node._attributes and attr in new_node._attributes \
           and hasattr(old_node, attr):
            setattr(new_node, attr, getattr(old_node, attr))
    return new_node


def fix_missing_locations(node):
    """
    When you compile a node tree with compile(), the compiler expects lineno
    and col_offset attributes for every node that supports them.  This is
    rather tedious to fill in for generated nodes, so this helper adds these
    attributes recursively where not already set, by setting them to the values
    of the parent node.  It works recursively starting at *node*.
    """
    def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
        if 'lineno' in node._attributes:
            if not hasattr(node, 'lineno'):
                node.lineno = lineno
            else:
                lineno = node.lineno
        if 'end_lineno' in node._attributes:
            if not hasattr(node, 'end_lineno'):
                node.end_lineno = end_lineno
            else:
                end_lineno = node.end_lineno
        if 'col_offset' in node._attributes:
            if not hasattr(node, 'col_offset'):
                node.col_offset = col_offset
            else:
                col_offset = node.col_offset
        if 'end_col_offset' in node._attributes:
            if not hasattr(node, 'end_col_offset'):
                node.end_col_offset = end_col_offset
            else:
                end_col_offset = node.end_col_offset
        for child in iter_child_nodes(node):
            _fix(child, lineno, col_offset, end_lineno, end_col_offset)
    _fix(node, 1, 0, 1, 0)
    return node


def increment_lineno(node, n=1):
    """
    Increment the line number and end line number of each node in the tree
    starting at *node* by *n*. This is useful to "move code" to a different
    location in a file.
    """
    for child in walk(node):
        if 'lineno' in child._attributes:
            child.lineno = (getattr(child, 'lineno', 0) or 0) + n
        if 'end_lineno' in child._attributes:
            child.end_lineno = (getattr(child, 'end_lineno', 0) or 0) + n
    return node