import ast
import inspect
import warnings
import linecache
from uuid import uuid4
from textwrap import dedent
from collections import OrderedDict
from astmonkey.transformers import ParentNodeTransformer
from astmonkey.visitors import SourceGeneratorNodeVisitor


SPLIT = 'split'


def split(func, name=None):
    return func, name


def value(val, name=None):
    return val, name


def get_sub(lst, pred):
    for subnode in lst:
        if pred(subnode):
            lst.remove(subnode)
            return subnode


class ToSource(SourceGeneratorNodeVisitor):
    '''
    The changes in here are primarily to patch
    around differences between python versions
    (I've tested with Python 3.5)
    '''

    @classmethod
    def to_source(cls, node):
        generator = cls('    ')
        generator.visit(node)
        return ''.join(generator.result)

    def visit_Call(self, node):
        node.starargs = get_sub(
            node.args,
            lambda t: isinstance(t, ast.Starred)
        )
        node.kwargs = get_sub(
            node.keywords,
            lambda t: t.arg is None
        )
        return super().visit_Call(node)

    def visit_Lambda(self, node):
        # until i can get the ast compiling directly,
        # we need to make sure precendence is correct
        self.write('(')
        super().visit_Lambda(node)
        self.write(')')

    def signature(self, node):
        if node.vararg:
            node.vararg = node.vararg.arg
        if node.kwarg:
            node.kwarg = node.kwarg.arg
        return super().signature(node)


class Visitor(ast.NodeVisitor):

    def __init__(self):
        self.parts = {}
        self.current = None
        self.name = None
        self.last_idx = -1

    def visit_FunctionDef(self, node):
        self.current = self.parts[node] = OrderedDict()
        self.name = node.name
        self.args = node.args
        return super().generic_visit(node)

    def function_from(self, name, args, body, lineno):
        func = ast.FunctionDef(
            name=name,
            args=args,
            body=body,
            decorator_list=[],
            returns=None
        )
        func.lineno = lineno
        return func

    def visit_Call(self, node):
        if getattr(node.func, 'id', None) == SPLIT:
            # this is the expression that contains the call,
            # or basically the value of the assignment/return
            expr = node.parent
            user = expr.parent  # by user we mean user of the value

            # it's possible more are actually supported,
            # but i'm hesitant to just allow them without
            # further testing
            if not isinstance(user, (ast.Return, ast.Assign)):
                warnings.warn('This is untested')

            # sanity checking
            assert user.parent_field == 'body', user.parent_field
            assert isinstance(user.parent, ast.FunctionDef)

            field = getattr(user.parent, user.parent_field)
            body = field[self.last_idx + 1:user.parent_field_index]

            value = user.value
            value.func = value.func.args[0]  # remove call to split
            body.append(ast.Return(value))
            body[-1].lineno = value.lineno

            self.current[self.name] = self.function_from(
                self.name, self.args, body, value.lineno
            )

            self.last_idx = user.parent_field_index

            name = node.args[1].s if len(node.args) == 2 else None
            self.name = name
            self.args = self.get_args(user, name)

        return super().generic_visit(node)

    def get_args(self, user, name):
        msg = 'have a name, on line {} of function {}'.format(
            user.lineno,
            user.parent.name
        )

        if isinstance(user, ast.Return):
            if name is not None:
                raise Exception('A closing split cannot {}'.format(msg))
        else:
            if name is None:
                raise Exception('A non-closing split must {}'.format(msg))

        args = []
        if isinstance(user, ast.Assign):
            target = user.targets[0]
            names = target.elts if isinstance(target, ast.Tuple) else [target]
            args = [ast.arg(name.id, None) for name in names]

        return ast.arguments(
            args=args,
            vararg=None,
            kwonlyargs=[],
            kw_defaults=[],
            kwarg=None,
            defaults=[]
        )


def cache_code(node):
    filename = '<ast_{}>'.format(uuid4().hex)
    source = ToSource.to_source(node)

    lines = [line + '\n' for line in source.splitlines()]
    linecache.cache[filename] = len(source), None, lines, filename
    assert filename in linecache.cache

    return filename


def extract(env, node, name):
    loc = dict(env)
    filename = cache_code(node)

    node = ast.Module(body=[node])
    node = ast.fix_missing_locations(node)

    code = compile(node, filename=filename, mode='exec')
    exec(code, loc, loc)
    return loc[name]


def rebuild(function):
    assert callable(function)
    lines, lineno = inspect.getsourcelines(function)

    lines = ''.join(lines)
    lines = dedent(lines)
    lines = '\n' * (lineno - 1) + lines

    root, = ast.parse(lines).body
    root = ParentNodeTransformer().visit(root)

    visitor = Visitor()
    visitor.visit(root)
    parts = visitor.parts[root]

    return OrderedDict(
        (name, extract(function.__globals__, node, name))
        for name, node in parts.items()
    )

@@ -1640,20 +1640,24 @@
    self.
-part
+function
 s = %7B%7D%0A 
@@ -3262,23 +3262,25 @@
    self.
-current
+functions
 %5Bself.na
@@ -5377,40 +5377,8 @@
 oot)
-%0A    parts = visitor.parts%5Broot%5D
 %0A%0A  
@@ -5483,20 +5483,32 @@
 node in 
-part
+visitor.function
 s.items(

Rename field