import ast
import inspect

from brian2.utils.stringtools import deindent, get_identifiers, indent

from .rendering import NodeRenderer

__all__ = [
    "AbstractCodeFunction",
    "abstract_code_from_function",
    "extract_abstract_code_functions",
    "substitute_abstract_code_functions",
]


class AbstractCodeFunction:
    """
    The information defining an abstract code function

    Has attributes corresponding to initialisation parameters

    Parameters
    ----------

    name : str
        The function name.
    args : list of str
        The arguments to the function.
    code : str
        The abstract code string consisting of the body of the function less
        the return statement.
    return_expr : str or None
        The expression returned, or None if there is nothing returned.
    """

    def __init__(self, name, args, code, return_expr):
        self.name = name
        self.args = args
        self.code = code
        self.return_expr = return_expr

    def __str__(self):
        s = (
            f"def {self.name}({', '.join(self.args)}):\n{indent(self.code)}\n    return"
            f" {self.return_expr}\n"
        )
        return s

    __repr__ = __str__


def abstract_code_from_function(func):
    """
    Converts the body of the function to abstract code

    Parameters
    ----------
    func : function, str or ast.FunctionDef
        The function object to convert. Note that the arguments to the
        function are ignored.

    Returns
    -------
    func : AbstractCodeFunction
        The corresponding abstract code function

    Raises
    ------
    SyntaxError
        If unsupported features are used such as if statements or indexing.
    """
    if callable(func):
        code = deindent(inspect.getsource(func))
        funcnode = ast.parse(code, mode="exec").body[0]
    elif isinstance(func, str):
        funcnode = ast.parse(func, mode="exec").body[0]
    elif func.__class__ is ast.FunctionDef:
        funcnode = func
    else:
        raise TypeError("Unsupported function type")

    if funcnode.args.vararg is not None:
        raise SyntaxError("No support for variable number of arguments")
    if funcnode.args.kwarg is not None:
        raise SyntaxError("No support for arbitrary keyword arguments")
    if len(funcnode.args.defaults):
        raise SyntaxError("No support for default values in functions")

    nodes = funcnode.body
    nr = NodeRenderer()
    lines = []
    return_expr = None
    for node in nodes:
        if node.__class__ is ast.Return:
            return_expr = nr.render_node(node.value)
            break
        else:
            lines.append(nr.render_node(node))
    abstract_code = "\n".join(lines)
    args = [arg.arg for arg in funcnode.args.args]
    name = funcnode.name
    return AbstractCodeFunction(name, args, abstract_code, return_expr)


def extract_abstract_code_functions(code):
    """
    Returns a set of abstract code functions from function definitions.

    Returns all functions defined at the top level and ignores any other
    code in the string.

    Parameters
    ----------
    code : str
        The code string defining some functions.

    Returns
    -------
    funcs : dict
        A mapping ``(name, func)`` for ``func`` an `AbstractCodeFunction`.
    """
    code = deindent(code)
    nodes = ast.parse(code, mode="exec").body
    funcs = {}
    for node in nodes:
        if node.__class__ is ast.FunctionDef:
            func = abstract_code_from_function(node)
            funcs[func.name] = func
    return funcs


class VarRewriter(ast.NodeTransformer):
    """
    Rewrites all variable names in names by prepending pre
    """

    def __init__(self, pre):
        self.pre = pre

    def visit_Name(self, node):
        return ast.Name(id=self.pre + node.id, ctx=node.ctx)

    def visit_Call(self, node):
        args = [self.visit(arg) for arg in node.args]
        return ast.Call(
            func=ast.Name(id=node.func.id, ctx=ast.Load()),
            args=args,
            keywords=[],
        )


class FunctionRewriter(ast.NodeTransformer):
    """
    Inlines a function call using temporary variables

    numcalls is the number of times the function rewriter has been called so
    far, this is used to make sure that when recursively inlining there is no
    name aliasing. The substitute_abstract_code_functions ensures that this is
    kept up to date between recursive runs.

    The pre attribute is the set of lines to be inserted above the currently
    being processed line, i.e. the inline code.

    The visit method returns the current line processed so that the function
    call is replaced with the output of the inlining.
    """

    def __init__(self, func, numcalls=0):
        self.func = func
        self.numcalls = numcalls
        self.pre = []
        self.suspend = False

    def visit_Call(self, node):
        # we suspend operations during an inlining operation, then resume
        # afterwards, see below, so we only ever try to expand one inline
        # function call at a time, i.e. no f(f(x)). This case is handled
        # by the recursion.
        if self.suspend:
            return node
        # We only work with the function we're provided
        if node.func.id != self.func.name:
            return node
        # Suspend while processing arguments (no recursion)
        self.suspend = True
        args = [self.visit(arg) for arg in node.args]
        self.suspend = False
        # The basename is used for function-local variables
        basename = f"_inline_{self.func.name}_{str(self.numcalls)}"
        # Assign all the function-local variables
        for argname, arg in zip(self.func.args, args):
            newpre = ast.Assign(
                targets=[ast.Name(id=f"{basename}_{argname}", ctx=ast.Store())],
                value=arg,
            )
            self.pre.append(newpre)
        # Rewrite the lines of code of the function using the names defined
        # above
        vr = VarRewriter(f"{basename}_")
        for funcline in ast.parse(self.func.code).body:
            self.pre.append(vr.visit(funcline))
        # And rewrite the return expression
        return_expr = vr.visit(ast.parse(self.func.return_expr, mode="eval").body)
        self.pre.append(
            ast.Assign(
                targets=[ast.Name(id=basename, ctx=ast.Store())], value=return_expr
            )
        )
        # Finally we replace the function call with the output of the inlining
        newnode = ast.Name(id=basename)
        self.numcalls += 1
        return newnode


def substitute_abstract_code_functions(code, funcs):
    """
    Performs inline substitution of all the functions in the code

    Parameters
    ----------
    code : str
        The abstract code to make inline substitutions into.
    funcs : list, dict or set of AbstractCodeFunction
        The function substitutions to use, note in the case of a dict, the
        keys are ignored and the function name is used.

    Returns
    -------
    code : str
        The code with inline substitutions performed.
    """
    if isinstance(funcs, (list, set)):
        newfuncs = dict()
        for f in funcs:
            newfuncs[f.name] = f
        funcs = newfuncs

    code = deindent(code)
    lines = ast.parse(code, mode="exec").body

    # This is a slightly nasty hack, but basically we just check by looking at
    # the existing identifiers how many inline operations have already been
    # performed by previous calls to this function
    ids = get_identifiers(code)
    funcstarts = {}
    for func in funcs.values():
        subids = {id for id in ids if id.startswith(f"_inline_{func.name}_")}
        subids = {id.replace(f"_inline_{func.name}_", "") for id in subids}
        alli = []
        for subid in subids:
            p = subid.find("_")
            if p > 0:
                subid = subid[:p]
            i = int(subid)
            alli.append(i)
        if len(alli) == 0:
            i = 0
        else:
            i = max(alli) + 1
        funcstarts[func.name] = i

    # Now we rewrite all the lines, replacing each line with a sequence of
    # lines performing the inlining
    newlines = []
    for line in lines:
        for func in funcs.values():
            rw = FunctionRewriter(func, funcstarts[func.name])
            line = rw.visit(line)
            newlines.extend(rw.pre)
            funcstarts[func.name] = rw.numcalls
        newlines.append(line)

    # Now we render to a code string
    nr = NodeRenderer()
    newcode = "\n".join(nr.render_node(line) for line in newlines)

    # We recurse until no changes in the code to ensure that all functions
    # are expanded if one function refers to another, etc.
    if newcode == code:
        return newcode
    else:
        return substitute_abstract_code_functions(newcode, funcs)