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.builder = ir.Builder()
      m = ir.Module(pkg.name)
      self.builder.setModule(m)
      self.genModule(pkg)
      return m

   # inner helpers:
   def genModule(self, pkg):
      for s in pkg.scope:
         if type(s) is astnodes.Variable:
            # TODO
            pass
         elif type(s) is astnodes.Function:
            # TODO: handle arguments
            # TODO handle return?
            bb = self.builder.newBB()
            self.builder.setBB(bb)
            self.genCode(s.body)
            self.builder.addIns(ir.Return())
         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)
         self.builder.addIns(ir.Store(code.lval, 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:
         pass
      elif type(code) is astnodes.EmptyStatement:
         pass
      elif type(code) is astnodes.ReturnStatement:
         pass
      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()
            print('Unknown cond', expr)
      elif type(expr) is astnodes.Literal:
         if expr.val:
            self.builder.addIns(ir.BranchInstruction(bbtrue))
         else:
            self.builder.addIns(ir.BranchInstruction(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)
         tmp = self.builder.newTmp()
         ops = ['+', '-', '*', '/', 'and', 'or']
         if expr.op in ops:
            op = expr.op
            ins = ir.BinaryOperator(tmp, op, ra, rb)
            self.builder.addIns(ins)
            return tmp
         else:
            print('Unknown {0}'.format(expr))
            # 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()
         i = ir.Load(expr, 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
      else:
         print('Unknown expr:', code)