from typing import Dict

from antlr4 import TerminalNode
from llvmlite import ir

from gen import ANFVisitor
from gen.ANFParser import ANFParser


class ExpBuilder(ANFVisitor):
    module = ir.Module("Program")
    i_type = ir.IntType(64)

    def __init__(self):
        self.var_env: Dict[ir.Value] = dict()
        self.currentFunc: ir.Function | None = None

    def visitDef(self, ctx: ANFParser.DefContext):
        args = [i.getText() for i in ctx.IDENT()][1:]
        self.currentFunc = ir.Function(self.module,
                                    ir.FunctionType(self.i_type, [self.i_type] * len(args)),
                                    name=ctx.IDENT(0))
        self.visit(ctx.cexp())


    def visitTrue(self, ctx:ANFParser.TrueContext):
        return ir.Constant(self.i_type, True)

    def visitFalse(self, ctx: ANFParser.FalseContext):
        return ir.Constant(self.i_type, False)

    def visitVar(self, ctx:ANFParser.VarContext):
        var_name = ctx.IDENT().getText()
        return self.var_env[var_name]

    def visitNum(self, ctx:ANFParser.NumContext):
        number = int(ctx.NUMBER().getText())


    def visitLet(self, ctx:ANFParser.LetContext):
        var_name = ctx.IDENT().getText()
        self.var_env[var_name] = self.visit(ctx.funcall())
        self.currentFunc.append_basic_block(self.var_env[var_name])
        return self.visit(ctx.cexp())