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

from .nodes import *
from ...core.errors import Error
from ... import core
from .builtin import real, integer, boolean, char

class KsIrGenerator:
   def __init__(self):
      pass

   # Code generation functions:
   def genexprcode(self, node):
      """ 
         Generate code for expressions!
         Recursively evaluates, and ensures a register contains the answer.
         register is an integer register or a floating point reg
      """
      if isinstance(node, Binop):
         """ Handle a binary operation (two arguments) of some kind """
         self.genexprcode(node.a)
         self.genexprcode(node.b)

         if node.op == 'mod':
            assert(node.typ.isType(integer))
         elif node.op == 'div':
            assert(node.typ.isType(integer))
         elif node.op == '*':
            if node.typ.isType(integer):
               pass
         elif node.op == '+':
            if node.typ.isType(integer):
               pass
         elif node.op == '-':
            if node.typ.isType(integer):
               pass
         else:
            Error('Unknown Binop {0}'.format(node.op))

      elif type(node) is Unop:
         if node.op == 'INTTOREAL':
            self.genexprcode(node.a)
            node.reg = node.a.reg
            # TODO use 'FILD' instruction
            freg = 12
            code.append('Unop inttoreal TODO')
         elif node.op == 'ABS':
            if isType(node.typ, real):
               code = [0xD9, 0xE1] # st(0) = fabs st(0)
               Error('ABS error integer')
            elif isType(node.typ, integer):
               code = []
               Error('ABS error integer')
            else:
               Error('ABS error')
         else:
            Error('Unknown Unop {0}'.format(node.op))

      elif isinstance(node, Designator):
         # dereference, array index. Make sure that the result comes into a register
         if len(node.selectors) > 0:
            self.gencode(node) # Load the pointer into some register

            # Now we can access the object at location '[node.reg]':
            if node.typ.isType(integer):
               self.addCode( mov(node.reg, [node.reg, 0x0]) )
            else:
               Error('Only integer types implemented')
         else:
            # No selectors, load variable directly
            if node.obj.typ.isType(integer):
               if type(node.obj) is Constant:
                  self.genexprcode(node.obj)
                  node.reg = node.obj.reg
               else:
                  pass
                  # Get a register to store the integer value
            else:
               Error('Cannot load variable type {0}'.format(node.typ))

      elif isinstance(node, Relop):
         # Create a boolean from operands
         # TODO create an alternative for expressions used as conditions.
         self.genexprcode(node.a)
         self.genexprcode(node.b)

         if node.a.typ.isType(integer):
            self.freereg(node.b)
         else:
            Error('Relop not implemented for {0}'.format(node.a.typ))

      elif type(node) is Constant:
         print('TODO: constant')

      elif type(node) is ProcedureCall:
         Error('TODO: proc call')
      else:
         Error('Cannot generate expression code for: {0}'.format(node))

   def gencode(self, node):
      """ Code generation function for AST nodes """
      if isinstance(node, Module):
         # TODO: recurse!
         
         self.mod = core.Module()
         # Create a function called init for this module:
         ftype = core.FunctionType(self.context.VoidType, [])
         func = core.Function(ftype, "init", self.mod)
         bb = self.gencode(node.initcode)
         self.mod.dump()
         return self.mod

      elif type(node) is Procedure:
        # calculate offsets for local variables and parameters
        # Variable location relative to 'rbp' register
        variables = node.symtable.getAllLocal(Variable)

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

      elif type(node) is ProcedureCall:
         # Prepare parameters on the stack:
         assert(len(node.args) == len(node.proc.typ.parameters))
         for arg, param in zip(node.args, node.proc.typ.parameters):

            if param.kind == 'value': 
               self.genexprcode(arg)
               self.addCode( push(arg.reg) )
               self.freereg( arg )
               stacksize += 8
            else:
               Error('Parameter kind other than value')

      elif type(node) is Assignment:
         if node.lval.typ.isType(integer):
           # TODO if node.rval is Constant of some datatype, move it to mem directly
           self.genexprcode(node.rval) # Calculate the value that has to be stored.
           print("TODO")
         else:
            Error('Assignments of other types not implemented')
            # TODO if left and right are designators, do some sort of memcpy.

      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 type(node) is AsmCode:
         def processOperand(op):
            if type(op) is list:
               if type(op[0]) is Variable:
                  var = op[0]
                  if var.isLocal:
                     return ['rbp', var.offset]
                  else:
                     Error('Can only use local variables in inline assembler')
            return op
         for asmline in node.asmcode:
            opcode, operands = asmline
            operands = [processOperand(opx) for opx in operands]
            print('assembling', opcode, *operands)
            func,nargs = opcodes[opcode]
            code = func(*operands)
            self.addCode(code)

      elif isinstance(node, EmptyStatement):
         pass

      elif type(node) is StringConstant:
        self.strings.append(node)
        self.data.append(node.value) # Add string to the data section

      elif type(node) is Designator:
         if len(node.selectors) > 0:
            self.getreg(node)
            # Load starting address
            if node.obj.isLocal:
               self.addCode( leareg64(node.reg, ['rbp', node.obj.offset]) )
            else:
               # Global variables need to be relocated...
               self.addCode(leareg64(node.reg, ['RIP', 0]))
               self.fixCode(self.rip - 4, imm32(node.obj.offset - self.rip))
            # Loop over all designators..
            for selector in node.selectors:
               if type(selector) is Index:
                  # Deref an array index
                  self.genexprcode(selector.index)
                  self.freereg(selector)
               elif type(selector) is Field:
                  print('Field')
                  Error('Field not implemented')
               else:
                  Error('Unknown selector')
         else:
            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)