import ast
import inspect

from ast import (
    copy_location,

    Assign, Load, Name, Store, Tuple, While
)
from functools import wraps


def deindent(source):
    if source.startswith(" "):
        lines = source.split("\n")
        level = len(filter(lambda s: not s, lines[0].split(" ")))
        result = []
        for line in lines:
            result.append(line[level:])

        return "\n".join(result)

    return source


def isa(n, k):
    return isa_and_has(n, k)


def isa_and_has(n, k, **kwargs):
    if not isinstance(n, k):
        return False

    for key, value in kwargs.items():
        if getattr(n, key) != value:
            return False

    return True


class TCOTransformer(ast.NodeTransformer):
    def __init__(self, name):
        self.name = name
        self.args = None

    def visit_Return(self, node):
        if isa(node.value, ast.Call):
            call = node.value
            if isa_and_has(call.func, Name, id=self.name):
                value = Tuple(call.args, Load())
                targets = Tuple([], Store())
                for i, argument in enumerate(self.args.args):
                    targets.elts.append(Name(argument.id, ast.Store()))

                return copy_location(Assign([targets], value), node)

        if isa(node.value, Name):
            return node

        raise ValueError("invalid expression in tail position")

    def visit_FunctionDef(self, node):
        self.args = node.args

        decorators, body = [], []
        for child in node.decorator_list:
            if isinstance(child, Name) and child.id == "tco":
                continue

            decorators.append(child)

        for child in node.body:
            body.append(self.visit(child))

        while_block = While(ast.Num(1), [], [])
        while_block.body = body
        node.decorator_list = decorators
        node.body = [while_block]
        return node


def tco(f):
    name = f.__name__
    filename = inspect.getsourcefile(f)
    source = inspect.getsource(f)
    module = inspect.getmodule(f)
    tree = ast.parse(deindent(source))
    tree = TCOTransformer(name).visit(tree)
    tree = ast.fix_missing_locations(tree)
    code = compile(tree, filename, "exec")
    globals_, locals_ = module.__dict__, {}
    exec code in globals_, locals_
    return wraps(f)(locals_[name])

@@ -69,16 +69,28 @@
  Assign,
+ Dict, List,
  Load, N
@@ -97,15 +97,25 @@
 ame,
+ Num,
  Store,
+ Str,
  Tup
@@ -324,117 +324,53 @@
   re
-sult = %5B%5D%0A        for line in lines:%0A            result.append(line%5Blevel:%5D)%0A%0A        return %22%5Cn%22.join(result
+turn %22%5Cn%22.join(line%5Blevel:%5D for line in lines
 )%0A%0A 
@@ -1226,33 +1226,40 @@
 )%0A%0A        if is
-a
+instance
 (node.value, Nam
@@ -1255,20 +1255,51 @@
 .value, 
+(
 Name
+, Num, Str, List, Tuple, Dict)
 ):%0A     

allow more exps in tail position