import ir
from . import astnodes

def NumGen():
   a = 0
   while True:
      yield a
      a = a + 1

class NameGenerator:
   def __init__(self, prefix):
      self.prefix = prefix
      self.nums = NumGen()
   def gen(self):
      return '{0}{1}'.format(self.prefix, self.nums.__next__())

class CodeGenerator:
   """ Generates intermediate code from a package """
   def gencode(self, pkg):
      assert type(pkg) is astnodes.Package
      self.m = ir.Module(pkg.name)
      self.newTmp = NameGenerator('t').gen
      self.newLab = NameGenerator('lab').gen
      self.genModule(pkg)
      return self.m

   # Helpers:
   def addIns(self, i):
      self.m.Instructions.append(i)
   # 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?
            self.addIns(ir.LabelInstruction(s.name))
            self.genCode(s.body)
            self.addIns(ir.RetInstruction())
         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.addIns(ir.StoreInstruction(code.lval, re))
      elif type(code) is astnodes.IfStatement:
         cr = self.genExprCode(code.condition)
         t1, t2, te = self.newLab(), self.newLab(), self.newLab()
         self.addIns(ir.IfInstruction(cr, t1, t2))
         self.addIns(ir.LabelInstruction(t1))
         self.genCode(code.truestatement)
         self.addIns(ir.BranchInstruction(te))
         self.addIns(ir.LabelInstruction(t2))
         self.genCode(code.falsestatement)
         self.addIns(ir.LabelInstruction(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 genExprCode(self, expr):
      if type(expr) is astnodes.Binop:
         ra = self.genExprCode(expr.a)
         rb = self.genExprCode(expr.b)
         ops = ['+', '-', '*', '/', 'and', 'or']
         if expr.op in ops:
            op = expr.op
            tmp = self.newTmp()
            ins = ir.BinaryOperator(tmp, op, ra, rb)
            self.addIns(ins)
            return tmp
         else:
            print('Unknown binop {0}'.format(expr))
      elif type(expr) is astnodes.Constant:
         tmp = unique()
      elif type(expr) is astnodes.VariableUse:
         tmp = self.newTmp()
      elif type(expr) is astnodes.Literal:
         tmp = self.newTmp()
         ins = ir.MoveInstruction(tmp, expr.val)
         self.addIns(ins)
         return tmp
      else:
         print('Unknown expr:', code)