import ppci
import ir
import registerallocator

class AsmLabel:
   def __init__(self, lab):
      self.lab = lab
   def __repr__(self):
      return '{0}:'.format(self.lab)

class Op:
   def __init__(self, op, dst, src):
      self.op = op
      self.src = src
      self.dst = dst
   def __repr__(self):
      return '{0} {1}, {2}'.format(self.op, self.dst, self.src)

class Jmp:
   def __init__(self, j, target):
      self.j = j
      self.target = target
   def __repr__(self):
      return '{0} {1}'.format(self.j, self.target)

class X86CodeGen:
   def __init__(self, diag):
      self.diag = diag
      self.regs = ['rax', 'rbx', 'rcx', 'rdx']

   def emit(self, i):
      self.asm.append(i)

   def genBin(self, ir):
      self.asm = []
      # Allocate registers:
      ra = registerallocator.RegisterAllocator()
      # TODO: do not register allocate on intermediate code:
      ra.registerAllocate(ir, self.regs)
      self.genModule(ir)
      return self.asm

   def genModule(self, ir):
      for f in ir.Functions:
         self.genFunction(f)
   def genFunction(self, f):
      self.emit('global {0}'.format(f.name))
      self.emit(AsmLabel(f.name))
      self.emit(Jmp('jmp', f.entry.name))
      for bb in f.BasicBlocks:
         self.genBB(bb)
   def genBB(self, bb):
      self.emit(AsmLabel(bb.name))
      for i in bb.Instructions:
         self.genIns(i)
   def genIns(self, i):
      if type(i) is ir.BinaryOperator:
         ops = {'+':'add', '-':'sub', '*':'mul'}
         if i.operation in ops:
            self.emit(Op('mov', i.result.reg, i.value1.reg))
            self.emit(Op(ops[i.operation], i.result.reg, i.value2.reg))
         else:
            raise NotImplementedError('op {0}'.format(i.operation))
      elif type(i) is ir.Load:
         self.emit(Op('mov', i.value, '[{0}]'.format(i.location)))
      elif type(i) is ir.Return:
         self.emit('ret')
      elif type(i) is ir.Call:
         self.emit('call')
      elif type(i) is ir.ImmLoad:
         self.emit(Op('mov', i.target, i.value))
      elif type(i) is ir.Store:
         self.emit(Op('mov', '[{0}]'.format(i.location), i.value))
      elif type(i) is ir.ConditionalBranch:
         self.emit(Op('cmp', i.a, i.b))
         jmps = {'>':'jg', '<':'jl', '==':'je'}
         if i.cond in jmps:
            j = jmps[i.cond]
            self.emit(Jmp(j, i.lab1.name))
         else:
            raise NotImplementedError('condition {0}'.format(i.cond))
         self.emit(Jmp('jmp', i.lab2.name))
      elif type(i) is ir.Branch:
         self.emit(Jmp('jmp', i.target.name))
      elif type(i) is ir.Alloc:
         pass
      else:
         raise NotImplementedError('{0}'.format(i))