#------------------------------------------------------------------------------
# Copyright (c) 2013-2017, Nucleic Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#------------------------------------------------------------------------------
from contextlib import contextmanager

from atom.api import Atom, Bool, Int, List, Str

from ..compat import IS_PY3, USE_WORDCODE
from . import byteplay as bp


class CodeGenerator(Atom):
    """ A class for generating bytecode operations.

    """
    #: The free variables for the code.
    freevars = List()

    #: The arguments for the code.
    args = List()

    #: Number of kwonly arguments
    kwonlyargs = Int()

    #: Whether the code takes variadic args.
    varargs = Bool(False)

    #: Whether the code takes variadic kwargs.
    varkwargs = Bool(False)

    #: Whether the code object should get new locals.
    newlocals = Bool(False)

    #: The name for the code object.
    name = Str()

    #: The full name of the file which is being compiled.
    filename = Str()

    #: The first line number for the code object.
    firstlineno = Int(1)

    #: The docstring for the code object.
    docstring = Str()

    #: The list of generated byteplay code operations.
    code_ops = List()

    def to_code(self):
        """ Create a Python code object from the current code ops.

        """
        if not IS_PY3:
            bp_code = bp.Code(
                self.code_ops, self.freevars, self.args, self.varargs,
                self.varkwargs, self.newlocals, self.name, self.filename,
                self.firstlineno, self.docstring or None
                )
        else:
            bp_code = bp.Code(
                self.code_ops, self.freevars, self.args, self.kwonlyargs,
                self.varargs, self.varkwargs, self.newlocals, self.name,
                self.filename, self.firstlineno, self.docstring or None
                )
        return bp_code.to_code()

    def set_lineno(self, lineno):
        """ Set the current line number in the code.

        """
        self.code_ops.append(                           # TOS
            (bp.SetLineno, lineno),                     # TOS
        )

    def load_global(self, name):
        """ Load a global variable onto the TOS.

        """
        self.code_ops.append(                           # TOS
            (bp.LOAD_GLOBAL, name),                     # TOS -> value
        )

    def load_fast(self, name):
        """ Load a fast local variable onto the TOS.

        """
        self.code_ops.append(                           # TOS
            (bp.LOAD_FAST, name),                       # TOS -> value
        )

    def load_const(self, const):
        """ Load a const value onto the TOS.

        """
        self.code_ops.append(                           # TOS
            (bp.LOAD_CONST, const),                     # TOS -> value
        )

    def load_attr(self, name):
        """ Load an attribute from the object on TOS.

        """
        self.code_ops.append(                           # TOS -> obj
            (bp.LOAD_ATTR, name),                       # TOS -> value
        )

    def store_global(self, name):
        """ Store the TOS as a global.

        """
        self.code_ops.append(                           # TOS -> value
            (bp.STORE_GLOBAL, name),                    # TOS
        )

    def store_fast(self, name):
        """ Store the TOS as a fast local.

        """
        self.code_ops.append(                           # TOS -> value
            (bp.STORE_FAST, name),                      # TOS
        )

    def store_attr(self, name):
        """ Store the value at 2nd as an attr on 1st.

        """
        self.code_ops.append(                           # TOS -> val -> obj
            (bp.STORE_ATTR, name),                      # TOS
        )

    def delete_global(self, name):
        """ Delete a named global variable.

        """
        self.code_ops.append(                           # TOS
            (bp.DELETE_GLOBAL, name),                   # TOS
        )

    def delete_fast(self, name):
        """ Delete a named fast local variable.

        """
        self.code_ops.append(                           # TOS
            (bp.DELETE_FAST, name),                     # TOS
        )

    def return_value(self):
        """ Return the value from the TOS.

        """
        self.code_ops.append(                           # TOS -> value
            (bp.RETURN_VALUE, None),                    # TOS
        )

    def binary_subscr(self):
        """ Subscript the #2 item with the TOS.

        """
        self.code_ops.append(                           # TOS -> obj -> idx
            (bp.BINARY_SUBSCR, None),                   # TOS -> value
        )

    def binary_multiply(self):
        """ Multiple the 2 items on the TOS.

        """
        self.code_ops.append(                           # TOS -> val_1 -> val_2
            (bp.BINARY_MULTIPLY, None),                 # TOS -> retval
        )

    def binary_add(self):
        """ Multiple the 2 items on the TOS.

        """
        self.code_ops.append(                           # TOS -> val_1 -> val_2
            (bp.BINARY_ADD, None),                      # TOS -> retval
        )

    def dup_top(self):
        """ Duplicate the value on the TOS.

        """
        self.code_ops.append(                           # TOS -> value
            (bp.DUP_TOP, None),                         # TOS -> value -> value
        )

    def build_map(self, n=0):
        """ Build a map and store it onto the TOS.

        """
        self.code_ops.append(                           # TOS
            (bp.BUILD_MAP, n),                          # TOS -> map
        )

    def build_tuple(self, n=0):
        """ Build a tuple from items on the TOS.

        """
        if n == 0:
            self.code_ops.append(                       # TOS
                (bp.LOAD_CONST, ()),                    # TOS -> tuple
            )
        else:
            self.code_ops.append(                       # TOS
                (bp.BUILD_TUPLE, n),                    # TOS -> tuple
            )

    def build_list(self, n=0):
        """ Build a list from items on the TOS.

        """
        self.code_ops.append(                           # TOS
            (bp.BUILD_LIST, n),                         # TOS -> list
        )

    def store_map(self):
        """ Store the key/value pair on the TOS into the map at 3rd pos.

        """
        if IS_PY3:
            # On Python 3 emulates store_map using MAP_ADD
            # STORE_MAP was removed in Python 3.5
            self.code_ops.append(
                (bp.MAP_ADD, 1),
            )
        else:
            self.code_ops.append(                       # TOS -> map -> value -> key
                (bp.STORE_MAP, None),                   # TOS -> map
            )

    def store_subscr(self):
        """ Store the index/value pair on the TOS into the 3rd item.

        """
        self.code_ops.append(                           # TOS -> value -> obj -> index
            (bp.STORE_SUBSCR, None),                    # TOS
        )

    if not IS_PY3:
        def build_class(self):
            """ Build a class from the top 3 stack items.

            """
            self.code_ops.append(                           # TOS -> name -> bases -> dict
                (bp.BUILD_CLASS, None),                     # TOS -> class
            )

    else:
        def load_build_class(self):
            """ Build a class from the top 3 stack items.

            """
            self.code_ops.append(                           # TOS
                (bp.LOAD_BUILD_CLASS, None),                # TOS -> builtins.__build_class__
            )

    def make_function(self, n_defaults=0):
        """ Make a function from a code object on the TOS.

        """
        # qual_name is absent under Python 2
        self.code_ops.append(                           # TOS -> qual_name -> code -> defaults
            (bp.MAKE_FUNCTION, n_defaults),             # TOS -> func
        )

    def call_function(self, n_args=0, n_kwds=0):
        """ Call a function on the TOS with the given args and kwargs.

        """
        if USE_WORDCODE:
            if n_kwds:
                # kwargs_name should be a tuple listing the keyword
                # arguments names
                # TOS -> func -> args -> kwargs -> kwargs_names
                op, arg = bp.CALL_FUNCTION_KW, n_args+n_kwds
            else:
                op, arg = bp.CALL_FUNCTION, n_args
        else:
            argspec = ((n_kwds & 0xFF) << 8) + (n_args & 0xFF)
            op, arg = bp.CALL_FUNCTION, argspec  # TOS -> func -> args -> kwargs

        self.code_ops.append((op, arg))          # TOS -> retval

    def call_function_var(self, n_args=0, n_kwds=0):
        """ Call a var function on the TOS with the given args and kwargs.

        """
        if USE_WORDCODE:
            # Under Python 3.6 positional arguments should always be stored
            # in a tuple and keywords in a mapping.
            argspec = 1 if n_kwds else 0
        else:
            argspec = ((n_kwds & 0xFF) << 8) + (n_args & 0xFF)

        opcode = (bp.CALL_FUNCTION_EX if USE_WORDCODE else
                  bp.CALL_FUNCTION_VAR)
        self.code_ops.append(                           # TOS -> func -> args -> kwargs -> varargs
            (opcode, argspec),                          # TOS -> retval
        )

    def pop_top(self):
        """ Pop the value from the TOS.

        """
        self.code_ops.append(                           # TOS -> value
            (bp.POP_TOP, None),                         # TOS
        )

    def rot_two(self):
        """ Rotate the two values on the TOS.

        """
        self.code_ops.append(                           # TOS -> val_1 -> val_2
            (bp.ROT_TWO, None),                         # TOS -> val_2 -> val_1
        )

    def rot_three(self):
        """ Rotate the three values on the TOS.

        """
        self.code_ops.append(                           # TOS -> val_1 -> val_2 -> val_3
            (bp.ROT_THREE, None),                       # TOS -> val_3 -> val_1 -> val_2
        )

    def unpack_sequence(self, n):
        """ Unpack the sequence on the TOS.

        """
        self.code_ops.append(                           # TOS -> obj
            (bp.UNPACK_SEQUENCE, n),                    # TOS -> val_n -> val_2 -> val_1
        )

    @contextmanager
    def try_squash_raise(self):
        """ A context manager for squashing tracebacks.

        The code written during this context will be wrapped so that
        any exception raised will appear to have been generated from
        the code, rather than any function called by the code.

        """
        exc_label = bp.Label()
        end_label = bp.Label()
        self.code_ops.append(
            (bp.SETUP_EXCEPT, exc_label),               # TOS
        )
        yield
        self.code_ops.extend([                          # TOS
            (bp.POP_BLOCK, None),                       # TOS
            (bp.JUMP_FORWARD, end_label),               # TOS
            (exc_label, None),                          # TOS -> tb -> val -> exc
            (bp.ROT_THREE, None),                       # TOS -> exc -> tb -> val
            (bp.ROT_TWO, None),                         # TOS -> exc -> val -> tb
            (bp.POP_TOP, None),                         # TOS -> exc -> val
            (bp.RAISE_VARARGS, 2),                      # TOS
            (bp.JUMP_FORWARD, end_label),               # TOS
            (bp.END_FINALLY, None),                     # TOS
            (end_label, None),                          # TOS
        ])

    @contextmanager
    def for_loop(self, iter_var, fast_var=True):
        """ A context manager for creating for-loops.

        Parameters
        ----------
        iter_var : str
            The name of the loop iter variable.

        fast_var : bool, optional
            Whether the iter_var lives in fast locals. The default is
            True. If False, the iter_var is loaded from globals.

        """
        start_label = bp.Label()
        jump_label = bp.Label()
        end_label = bp.Label()
        load_op = bp.LOAD_FAST if fast_var else bp.LOAD_GLOBAL
        self.code_ops.extend([
            (bp.SETUP_LOOP, end_label),
            (load_op, iter_var),
            (bp.GET_ITER, None),
            (start_label, None),
            (bp.FOR_ITER, jump_label),
        ])
        yield
        self.code_ops.extend([
            (bp.JUMP_ABSOLUTE, start_label),
            (jump_label, None),
            (bp.POP_BLOCK, None),
            (end_label, None),
        ])

    def insert_python_block(self, pydata, trim=True):
        """ Insert the compiled code for a Python Module ast or string.

        """
        code = compile(pydata, self.filename, mode='exec')
        bp_code = bp.Code.from_code(code).code
        if trim:  # skip SetLineno and ReturnValue
            bp_code = bp_code[1:-2]
        self.code_ops.extend(bp_code)

    def insert_python_expr(self, pydata, trim=True):
        """ Insert the compiled code for a Python Expression ast or string.

        """
        code = compile(pydata, self.filename, mode='eval')
        bp_code = bp.Code.from_code(code).code
        if trim:  # skip ReturnValue
            bp_code = bp_code[:-1]
        self.code_ops.extend(bp_code)

    def rewrite_to_fast_locals(self, local_names):
        """ Rewrite the locals to be loaded from fast locals.

        Given a set of available local names, this method will rewrite
        the current code ops, replaces every instance of a *_NAME opcode
        with a *_FAST or *_GLOBAL depending on whether or not the name
        exists in local_names or was written via STORE_NAME. This method
        is useful to convert the code so it can be used as a function.

        Parameters
        ----------
        local_names : set
            The set of available locals for the code.

        Returns
        -------
        result : list
            The list of names which must be provided as arguments.

        """
        arg_names = []
        stored_names = set()
        code_ops = self.code_ops
        for idx, (op, op_arg) in enumerate(code_ops):
            if op == bp.STORE_NAME:
                stored_names.add(op_arg)
                code_ops[idx] = (bp.STORE_FAST, op_arg)
        for idx, (op, op_arg) in enumerate(code_ops):
            if op == bp.LOAD_NAME:
                if op_arg in local_names:
                    op = bp.LOAD_FAST
                    arg_names.append(op_arg)
                elif op_arg in stored_names:
                    op = bp.LOAD_FAST
                else:
                    op = bp.LOAD_GLOBAL
                code_ops[idx] = (op, op_arg)
            elif op == bp.DELETE_NAME:          # py2.6 list comps
                if op_arg in stored_names:
                    op = bp.DELETE_FAST
                else:
                    op = bp.DELETE_GLOBAL
                code_ops[idx] = (op, op_arg)
        self.args = arg_names
        self.newlocals = True
        return arg_names
