import ir
from . import astnodes

def genModule(pkg):
   m = ir.Module(pkg.name)
   for s in pkg.scope:
      if type(s) is astnodes.Variable:
         genGlobal(m, s)
      elif type(s) is astnodes.Function:
         genFunction(m, s)
      else:
         print(s)
   return m

def genGlobal(m, var):
   v = ir.Value()
   v.name = var.name
   m.Globals.append(v)

def genFunction(m, fnc):
   ft = genType(fnc.typ)
   f = ir.Function(fnc.name, ft)
   m.Functions.append(f)
   bb = ir.BasicBlock()
   f.BasicBlocks.append(bb)
   genCode(bb, fnc.body)
   bb.Instructions.append(ir.RetInstruction())

def genCode(bb, code):
   if type(code) is astnodes.CompoundStatement:
      for s in code.statements:
         genCode(bb, s)
   elif type(code) is astnodes.Assignment:
      genExprCode(bb, code.rval)
      # TODO: store
   elif type(code) is astnodes.IfStatement:
      genExprCode(bb, code.condition)
      # TODO: implement IF.
      t1, t2 = 1, 2
      b = ir.BranchInstruction(t1, t2)
      bb.Instructions.append(b)
      genCode(bb, code.truestatement)
      genCode(bb, code.falsestatement)
   elif type(code) is astnodes.FunctionCall:
      ins = ir.CallInstruction('f', [])
      bb.Instructions.append(ins)
   elif type(code) is astnodes.EmptyStatement:
      pass
   elif type(code) is astnodes.ReturnStatement:
      bb.Instructions.append(ir.RetInstruction())
   else:
      print('Unknown stmt:', code)

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

nums = NumGen()
def unique():
   return 'tmp{0}'.format(nums.__next__())

def genExprCode(bb, code):
   if type(code) is astnodes.Binop:
      a = genExprCode(bb, code.a)
      b = genExprCode(bb, code.b)
      ops = {'+': 'fadd', '-': 'fsub', '*':'fmul', '/':'fdiv'}
      if code.op in ops:
         op = ops[code.op]
         tmp = unique()
         ins = ir.BinaryOperator(tmp, op, a, b)
         bb.Instructions.append(ins)
         return tmp
      else:
         print('Unknown binop')
         bb.Instructions.append(ir.BinaryOperator('unk2', code.op, a, b))
         return 'unk2'
   elif type(code) is astnodes.Constant:
      tmp = unique()
      bb.Instructions.append(ir.LoadInstruction(tmp, code.value))
      return tmp
   elif type(code) is astnodes.VariableUse:
      tmp = unique()
      ins = ir.LoadInstruction(tmp, code.target.name)
      return tmp
   else:
      print('Unknown expr:', code)
      return 'unk'

def genType(t):
   return ir.Type()

class CodeGenerator:
   """ Generates intermediate code """
   def gencode(self, ast):
      assert type(ast) is astnodes.Package
      return genModule(ast)