import ast

from .wasm import Function, Import, Module, Statement

class Visitor(ast.NodeVisitor):
    def __init__(self, module: Module) -> None:
        self._stack = []
        self.module = module

    # def visit_ImportFrom(self, node):
    #     for alias in node.names:
    #         self.imports.append(Import(
    #             node.module,
    #             alias.name,
    #             alias.asname,
    #         ))

    def visit_FunctionDef(self, node):
        if node.decorator_list:
            # TODO: Support normal decorators
            assert 1 == len(node.decorator_list)
            call = node.decorator_list[0]
            assert 'external' == call.func.id
            assert 1 == len(call.args)
            assert isinstance(call.args[0].value, str)

            import_ = Import(
                call.args[0].value,
                node.name,
                node.name,
            )

            import_.params = [
                arg.annotation.id
                for arg in node.args.args
            ]

            self.module.imports.append(import_)
            return

        func = Function(
            node.name,
        )

        self._stack.append(func)
        self.generic_visit(node)
        self._stack.pop()

        self.module.functions.append(func)

    def visit_Expr(self, node):
        self.generic_visit(node)

    def visit_Call(self, node):
        self.generic_visit(node)

        func = self._stack[-1]
        func.statements.append(
            Statement('call', '$' + node.func.id)
        )

    def visit_BinOp(self, node):
        self.generic_visit(node)

        func = self._stack[-1]

        if 'Add' == node.op.__class__.__name__:
            func.statements.append(
                Statement('i32.add')
            )
        elif 'Mult' == node.op.__class__.__name__:
            func.statements.append(
                Statement('i32.mul')
            )
        else:
            raise NotImplementedError

    def visit_Constant(self, node):
        if not self._stack:
            # Constant outside of any function
            imp = self.imports[-1]
            prefix = imp.name + '('
            val = node.value.strip()

            if val.startswith(prefix) and val.endswith(')'):
                imp.params = val[len(prefix):-1].split(',')
        else:
            func = self._stack[-1]
            if isinstance(node.value, int):
                func.statements.append(
                    Statement('i32.const', str(node.value))
                )

        self.generic_visit(node)