import ir
from . import astnodes
from .scope import boolType
   
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.scope:
         if type(s) is astnodes.Function:
            f = self.builder.newFunction(s.name)
            self.funcMap[s] = f
      for s in pkg.scope:
         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.scope:
               #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):
      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)
         loc = self.varMap[code.lval.target]
         self.builder.addIns(ir.Store(loc, re))
      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.FunctionCall:
         print('TODO')
         pass
      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.dostatement)
         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 = {'+':'addtmp', '-':'subtmp', '*': 'multmp', '/':'divtmp'}
            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.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.Literal:
         tmp = self.builder.newTmp()
         ins = ir.ImmLoad(tmp, expr.val)
         self.builder.addIns(ins)
         return tmp
      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:
         print('Unknown expr:', code)