"""
  Generates ir code from ast tree.
"""

from .nodes import *
from .builtin import real, integer, boolean, char, void

def coreType(typ):
   """ Return the correct core type given a type """
   if type(typ) is BaseType:
      if typ is integer:
         return core.i32
      if typ is void:
         return core.void
   elif type(typ) is ProcedureType:
      rType = coreType(typ.returntype)
      fpTypes = [coreType(p.typ) for p in typ.parameters]
      return core.FunctionType(rType, fpTypes)
   print(typ)
   raise NotImplementedError()

class KsIrGenerator:
   def __init__(self):
      self.builder = core.IRBuilder()
   # Code generation functions:
   def genexprcode(self, node):
      """ Generate code for expressions! """
      if isinstance(node, Binop):
         """ Handle a binary operation (two arguments) of some kind """
         lhs = self.genexprcode(node.a)
         rhs = self.genexprcode(node.b)

         if node.op == '*':
            if node.typ.isType(integer):
               return self.builder.createMul(lhs, rhs)
         elif node.op == '+':
            if node.typ.isType(integer):
               return self.builder.createAdd(lhs, rhs)
         elif node.op == '-':
            if node.typ.isType(integer):
               return self.builder.createSub(lhs, rhs)
         Error('Unknown binop or type {0}'.format(node))

      elif isinstance(node, Designator):
         # dereference, array index. Make sure that the result comes into a register
         if len(node.selectors) > 0:
            Error('Only integer types implemented')
         else:
            # No selectors, load variable directly
            print(node)
            #Error('Cannot load variable type {0}'.format(node.typ))
      elif type(node) is Constant:
         return core.Constant(node.value, coreType(node.typ))
      else:
         Error('Cannot generate expression code for: {0}'.format(node))

   def gencode(self, node):
      """ Code generation function for AST nodes """
      if isinstance(node, Module):
         # Create module:
         self.mod = core.Module(node.name) 

         globs = node.symtable.getAllLocal(Variable)
         for g in globs:
            print('global:', g)
         # Loop over all functions:
         print(node.symtable)
         node.symtable.printTable()
         funcs = node.symtable.getAllLocal(Procedure)
         for f in funcs:
            self.gencode(f)
         # Create a function called init for this module:
         self.mod.dump()
         return self.mod

      elif type(node) is Procedure:
            ftype = coreType(node.typ)
            print('function', node, ftype)
            func = core.Function(ftype, node.name, self.mod)
            bb = core.BasicBlock()
            func.basicblocks.append(bb)
            self.builder.setInsertPoint(bb)
            self.gencode(node.block)
            self.builder.setInsertPoint(None)
            variables = node.symtable.getAllLocal(Variable)
            print(variables)

      elif isinstance(node, StatementSequence):
         for s in node.statements:
            self.gencode(s)

      elif type(node) is ProcedureCall:
         # Prepare parameters on the stack:
         print("TODO")

      elif type(node) is Assignment:
         if node.lval.typ.isType(integer):
           print('assign')
           rhs = self.genexprcode(node.rval) # Calculate the value that has to be stored.
           #self.gencode(node.lval)
           print("TODO: assigment")
           
         else:
            Error('Assignments of other types not implemented')

      elif type(node) is IfStatement:
        self.genexprcode(node.condition)
        print("TODO IF")
        if node.falsestatement:
           # If with else clause
           pass
        else:
           # If without else clause
           pass

      elif isinstance(node, WhileStatement):
        self.genexprcode(node.condition)
        self.gencode(node.dostatements)
      elif type(node) is ForStatement:
         # Initial load of iterator variable:
         self.genexprcode(node.begin)
         self.genexprcode(node.end)
         self.gencode(node.statements)
         Error('No implementation of FOR statement')

      elif isinstance(node, EmptyStatement):
         pass # That was easy :)

      elif type(node) is StringConstant:
        self.strings.append(node)

      elif type(node) is Designator:
         Error('Can only gencode for designator with selectors')
      else:
         print('not generating code for {0}'.format(node))

   def generateIr(self, context, ast):
     """ ir generation front end """
     # Create a new context for this code.
     self.context = context
     return self.gencode(ast)