import ast
from xonsh.built_ins import XSH

macros = {}
def custom_macro(f, macros = macros):
    macros[f.__qualname__] = f
    return f

class Chained_Lambda:
    F,S,T = s = range(3)

    def __init__(self, *chain):
        self.chain = list(chain) if chain else []

    def __or__(self, x):
        if isinstance(x, Chained_Lambda):
            return Chained_Lambda(*self.chain + x.chain)
        elif hasattr(x, '__iter__'):
            return Chained_Lambda(*self.chain + list(x))
        else:
            return Chained_Lambda(*self.chain + [x])

    def __call__(self, n, l=True):
        chain = []
        for i in self.chain:
            if i not in self.s and (len(chain) == 0 or chain[-1] not in self.s):
                chain.append(self.T)
            chain.append(i)

        for i, f in [chain[i:i+2] for i in range(0, len(chain), 2)]:
            if i == self.F:
                n = filter(f, n)
                continue
            if i == self.S:
                n = sorted(n, key=f)
                continue
            if i == self.T:
                n = map(f, n)
                continue
        return list(n) if l else n

def get_node_locs(node):
    return {
        'lineno': node.lineno,
        'col_offset': node.col_offset
    }

@custom_macro
def λ(node):
    locs = get_node_locs(node)
    
    keyargs = {i.arg: i.value for i in node.keywords} or {i: ast.Constant(None, **locs) for i in "xyzw"}
    return ast.Lambda(
        args=ast.arguments(
            posonlyargs=[],
            kwonlyargs=[],
            args=[ast.arg(arg=i, **locs) for i in keyargs.keys()],
            defaults=list(keyargs.values()),
            kw_defaults=[]
        ),
        body=node.args[0],
        **locs
    )

def tranformation_lambda(node, value):
    locs = get_node_locs(node)
    
    return ast.Call(
        func=ast.Name(
            id='Chained_Lambda',
            ctx=ast.Load(),
            **locs
        ),
        args=[
            ast.Constant(value=value, **locs),
            λ(node)
        ],
        keywords=[],
        **locs
    )

@custom_macro
def fλ(node): return tranformation_lambda(node, 0)
@custom_macro
def sλ(node): return tranformation_lambda(node, 1)
@custom_macro
def tλ(node): return tranformation_lambda(node, 2)

class Macro_transformer(ast.NodeTransformer):
    def __init__(self, macros):
        self.macros = macros
    
    def generic_visit(self, node):
        super().generic_visit(node)
        
        if isinstance(node, ast.Call) and hasattr(node, 'func') and hasattr(node.func, 'id') and node.func.id in self.macros:
            node = self.macros[node.func.id](node)
        return node

XSH.execer.parser.custom_tree_process = lambda tree: Macro_transformer(macros).visit(tree)