import ir
from . import astnodes
from .scope import boolType, intType
from ppci import CompilerError
   
class CodeGenerator:
    """ Generates intermediate code from a package """
    def gencode(self, pkg):
        assert type(pkg) is astnodes.Package
        self.varMap = {} # Maps variables to storage locations.
        self.funcMap = {}
        self.builder = ir.Builder()
        m = ir.Module(pkg.name)
        self.builder.setModule(m)
        self.genModule(pkg)
        return m

    # inner helpers:
    def genModule(self, pkg):
       # Take care of forward declarations:
       for s in pkg.innerScope.Functions:
            f = self.builder.newFunction(s.name)
            self.funcMap[s] = f
       for s in pkg.innerScope:
         if type(s) is astnodes.Variable:
              v = self.builder.newVariable(s.name)
              #self.builder.addIns(ir.Alloc(v))
              self.varMap[s] = v
         elif type(s) is astnodes.Function:
            # TODO: handle arguments
            f = self.funcMap[s]
            self.builder.setFunction(f)
            bb = self.builder.newBB()
            f.entry = bb
            self.builder.setBB(bb)
            # generate room for locals:

            for sym in s.innerScope:
               #print(sym, sym.isParameter)
               # TODO: handle parameters different
               v = self.builder.newTmp(sym.name)
               self.builder.addIns(ir.Alloc(v))
               self.varMap[sym] = v

            self.genCode(s.body)
            # TODO handle return?
            self.builder.addIns(ir.Return())
            self.builder.setFunction(None)
         else:
            print(s)

    def genCode(self, code):
        assert isinstance(code, astnodes.Statement)
        if type(code) is astnodes.CompoundStatement:
            for s in code.statements:
                self.genCode(s)
        elif type(code) is astnodes.Assignment:
         re = self.genExprCode(code.rval)
         # TODO: Handle pointers
         loc = self.genExprCode(code.lval)
         # determine location of variable
         self.builder.addIns(ir.Store(loc, re))
        elif type(code) is astnodes.ExpressionStatement:
            self.genExprCode(code.ex)
        elif type(code) is astnodes.IfStatement:
         bbtrue = self.builder.newBB()
         bbfalse = self.builder.newBB()
         te = self.builder.newBB()
         self.genCondCode(code.condition, bbtrue, bbfalse)
         self.builder.setBB(bbtrue)
         self.genCode(code.truestatement)
         self.builder.addIns(ir.Branch(te))
         self.builder.setBB(bbfalse)
         self.genCode(code.falsestatement)
         self.builder.addIns(ir.Branch(te))
         self.builder.setBB(te)
        elif type(code) is astnodes.EmptyStatement:
         pass
        elif type(code) is astnodes.ReturnStatement:
         if code.expr:
            re = self.genExprCode(code.expr)
            self.builder.addIns(ir.Return(re))
         else:
            self.builder.addIns(ir.Return())
        elif type(code) is astnodes.WhileStatement:
         bbdo = self.builder.newBB()
         bbtest = self.builder.newBB()
         te = self.builder.newBB()
         self.builder.addIns(ir.Branch(bbtest))
         self.builder.setBB(bbtest)
         self.genCondCode(code.condition, bbdo, te)
         self.builder.setBB(bbdo)
         self.genCode(code.statement)
         self.builder.addIns(ir.Branch(bbtest))
         self.builder.setBB(te)
        else:
         print('Unknown stmt:', code)
    def genCondCode(self, expr, bbtrue, bbfalse):
      # Implement sequential logical operators
      assert expr.typ == boolType
      if type(expr) is astnodes.Binop:
         if expr.op == 'or':
            l2 = self.builder.newBB()
            self.genCondCode(expr.a, bbtrue, l2)
            self.builder.setBB(l2)
            self.genCondCode(expr.b, bbtrue, bbfalse)
         elif expr.op == 'and':
            l2 = self.builder.newBB()
            self.genCondCode(expr.a, l2, bbfalse)
            self.builder.setBB(l2)
            self.genCondCode(expr.b, bbtrue, bbfalse)
         elif expr.op in ['==', '>', '<']:
            ta = self.genExprCode(expr.a)
            tb = self.genExprCode(expr.b)
            i = ir.ConditionalBranch(ta, expr.op, tb, bbtrue, bbfalse)
            self.builder.addIns(i)
         else:
            raise NotImlementedError('Unknown condition {0}'.format(expr))
      elif type(expr) is astnodes.Literal:
         if expr.val:
            self.builder.addIns(ir.Branch(bbtrue))
         else:
            self.builder.addIns(ir.Branch(bbfalse))
      else:
         print('Unknown cond', expr)
    def genExprCode(self, expr):
        if type(expr) is astnodes.Binop:
         ra = self.genExprCode(expr.a)
         rb = self.genExprCode(expr.b)
         ops = ['+', '-', '*', '/', '|', '&']
         if expr.op in ops:
            tmpnames = {'+':'add', '-':'sub', '*': 'mul', '/':'div', '|':'or', '&':'and'}
            tmp = self.builder.newTmp(tmpnames[expr.op])
            op = expr.op
            ins = ir.BinaryOperator(tmp, op, ra, rb)
            self.builder.addIns(ins)
            return tmp
         else:
            print('Unknown {0}'.format(expr))
            tmp = self.builder.newTmp()
            # TODO
            return tmp
        elif type(expr) is astnodes.Unop:
            ra = self.genExprCode(expr.a)
            if expr.op == '&':
                # Address of operator
                tmp = self.builder.newTmp('addr')
                return tmp
                #self.builder.addIns(ins)
            else:
                raise Exception('Unknown {0}'.format(expr))
        elif type(expr) is astnodes.Constant:
            tmp = self.builder.newTmp()
            # TODO
            return tmp
        elif type(expr) is astnodes.VariableUse:
            tmp = self.builder.newTmp()
            loc = self.varMap[expr.target]
            i = ir.Load(loc, tmp)
            self.builder.addIns(i)
            return tmp
        elif type(expr) is astnodes.Deref:
            # dereference pointer type:
            addr = self.genExprCode(expr.ptr)
            tmp = self.builder.newTmp('deref')
            ins = ir.Load(addr, tmp)
            self.builder.addIns(ins)
            return tmp
        elif type(expr) is astnodes.Literal:
            tmp = self.builder.newTmp()
            ins = ir.ImmLoad(tmp, expr.val)
            self.builder.addIns(ins)
            return tmp
        elif type(expr) is astnodes.TypeCast:
            ar = self.genExprCode(expr.a)
            if isinstance(expr.to_type, astnodes.PointerType):
                if expr.a.typ is intType:
                    return ar
                elif isinstance(expr.a.typ, astnodes.PointerType):
                    return ar
                else:
                    raise Exception()
        elif type(expr) is astnodes.FunctionCall:
         tmp = self.builder.newTmp("res")
         args = []
         for arg in expr.args:
            ar = self.genExprCode(arg)
            args.append(ar)
         fn = self.funcMap[expr.proc]
         ins = ir.Call(fn, args, tmp)
         self.builder.addIns(ins)
         return tmp
        else:
         raise CompilerError('Unknown expr {}'.format(expr))