import ir
from . import astnodes
from .scope import boolType, intType
from ppci import CompilerError
from .typecheck import theType

tmpnames = {'+':'add', '-':'sub', '*': 'mul', '/':'div', '|':'or', \
    '&':'and', '>>':'shl', '<<':'shr'}

class CodeGenerator:
    """ Generates intermediate code from a package """
    def gencode(self, pkg):
        assert type(pkg) is astnodes.Package
        self.varMap = {} # Maps variables to storage locations.
        self.funcMap = {}
        self.builder = ir.Builder()
        m = ir.Module(pkg.name)
        self.builder.setModule(m)
        self.genModule(pkg)
        return m

    # inner helpers:
    def genModule(self, pkg):
       # Take care of forward declarations:
       for s in pkg.innerScope.Functions:
            f = self.builder.newFunction(s.name)
            self.funcMap[s] = f
       for s in pkg.innerScope:
         if type(s) is astnodes.Variable:
              v = self.builder.newTmp(s.name)
              #self.builder.addIns(ir.Alloc(v))
              self.varMap[s] = v
         elif type(s) is astnodes.Function:
            # TODO: handle arguments
            f = self.funcMap[s]
            self.builder.setFunction(f)
            bb = self.builder.newBB()
            f.entry = bb
            self.builder.setBB(bb)
            # generate room for locals:

            for sym in s.innerScope:
               #print(sym, sym.isParameter)
               # TODO: handle parameters different
               v = self.builder.newTmp(sym.name)
               self.builder.addIns(ir.Alloc(v))
               self.varMap[sym] = v

            self.genCode(s.body)
            # TODO handle return?
            self.builder.addIns(ir.Return())
            self.builder.setFunction(None)
         else:
            print(s)

    def genCode(self, code):
        assert isinstance(code, astnodes.Statement)
        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)
            loc = self.genExprCode(code.lval)
            self.builder.addIns(ir.Store(loc, re))
        elif type(code) is astnodes.ExpressionStatement:
            self.genExprCode(code.ex)
        elif type(code) is astnodes.IfStatement:
             bbtrue = self.builder.newBB()
             bbfalse = self.builder.newBB()
             te = self.builder.newBB()
             self.genCondCode(code.condition, bbtrue, bbfalse)
             self.builder.setBB(bbtrue)
             self.genCode(code.truestatement)
             self.builder.addIns(ir.Branch(te))
             self.builder.setBB(bbfalse)
             if code.falsestatement:
                 self.genCode(code.falsestatement)
             self.builder.addIns(ir.Branch(te))
             self.builder.setBB(te)
        elif type(code) is astnodes.ReturnStatement:
            if code.expr:
                re = self.genExprCode(code.expr)
                self.builder.addIns(ir.Return(re))
            else:
                self.builder.addIns(ir.Return())
        elif type(code) is astnodes.WhileStatement:
            bbdo = self.builder.newBB()
            bbtest = self.builder.newBB()
            te = self.builder.newBB()
            self.builder.addIns(ir.Branch(bbtest))
            self.builder.setBB(bbtest)
            self.genCondCode(code.condition, bbdo, te)
            self.builder.setBB(bbdo)
            self.genCode(code.statement)
            self.builder.addIns(ir.Branch(bbtest))
            self.builder.setBB(te)
        else:
            print('Unknown stmt:', code)

    def genCondCode(self, expr, bbtrue, bbfalse):
        # Implement sequential logical operators
        assert expr.typ == boolType
        if type(expr) is astnodes.Binop:
             if expr.op == 'or':
                l2 = self.builder.newBB()
                self.genCondCode(expr.a, bbtrue, l2)
                self.builder.setBB(l2)
                self.genCondCode(expr.b, bbtrue, bbfalse)
             elif expr.op == 'and':
                l2 = self.builder.newBB()
                self.genCondCode(expr.a, l2, bbfalse)
                self.builder.setBB(l2)
                self.genCondCode(expr.b, bbtrue, bbfalse)
             elif expr.op in ['==', '>', '<']:
                ta = self.genExprCode(expr.a)
                tb = self.genExprCode(expr.b)
                i = ir.ConditionalBranch(ta, expr.op, tb, bbtrue, bbfalse)
                self.builder.addIns(i)
             else:
                raise NotImlementedError('Unknown condition {0}'.format(expr))
        elif type(expr) is astnodes.Literal:
             if expr.val:
                self.builder.addIns(ir.Branch(bbtrue))
             else:
                self.builder.addIns(ir.Branch(bbfalse))
        else:
             print('Unknown cond', expr)

    def cast_to_rvalue(self, expr, loc):
        """ Cast lvalue to rvalue if required """
        if hasattr(expr, 'expect_rvalue'):
            # Load value from memory location:
            tmp = self.builder.newTmp()
            i = ir.Load(loc, tmp)
            self.builder.addIns(i)
            return tmp
        if expr.lvalue:
            return loc

    def genExprCode(self, expr):
        assert isinstance(expr, astnodes.Expression)
        if type(expr) is astnodes.Binop:
             ra = self.genExprCode(expr.a)
             rb = self.genExprCode(expr.b)
             ops = ['+', '-', '*', '/', '|', '&', '<<', '>>']
             if expr.op in ops:
                tmp = self.builder.newTmp(tmpnames[expr.op])
                op = expr.op
                ins = ir.BinaryOperator(tmp, op, ra, rb)
                self.builder.addIns(ins)
                return tmp
             else:
                raise NotImplementedError('Unknown {0}'.format(expr))
                tmp = self.builder.newTmp()
                # TODO
                return tmp
        elif type(expr) is astnodes.Unop:
            ra = self.genExprCode(expr.a)
            if expr.op == '&':
                # Address of operator
                tmp = self.builder.newTmp('addr')
                return tmp
                #self.builder.addIns(ins)
            else:
                raise NotImplementedError('Unknown {0}'.format(expr))
        elif type(expr) is astnodes.Constant:
            tmp = self.builder.newTmp()
            # TODO
            return tmp
        elif type(expr) is astnodes.VariableUse:
            loc = self.varMap[expr.target]
            return self.cast_to_rvalue(expr, loc)
        elif type(expr) is astnodes.Deref:
            # dereference pointer type:
            addr = self.genExprCode(expr.ptr)
            tmp = self.builder.newTmp('deref')
            ins = ir.Load(addr, tmp)
            self.builder.addIns(ins)
            return tmp
        elif type(expr) is astnodes.FieldRef:
            b = self.genExprCode(expr.base)
            offset = self.builder.newTmp('off_' + expr.field)
            bt = theType(expr.base.typ)
            ofs = bt.fieldOffset(expr.field)
            ins = ir.ImmLoad(offset, ofs)
            self.builder.addIns(ins)
            tmp2 = self.builder.newTmp('adr_' + expr.field)
            ins = ir.BinaryOperator(tmp2, '+', b, offset)
            self.builder.addIns(ins)
            return self.cast_to_rvalue(expr, tmp2)
        elif type(expr) is astnodes.Literal:
            tmp = self.builder.newTmp()
            ins = ir.ImmLoad(tmp, expr.val)
            self.builder.addIns(ins)
            return tmp
        elif type(expr) is astnodes.TypeCast:
            ar = self.genExprCode(expr.a)
            tt = theType(expr.to_type)
            if isinstance(tt, astnodes.PointerType):
                if expr.a.typ is intType:
                    return ar
                elif isinstance(expr.a.typ, astnodes.PointerType):
                    return ar
                else:
                    raise Exception()
            else:
                raise NotImplementedError("not implemented")
        elif type(expr) is astnodes.FunctionCall:
             tmp = self.builder.newTmp("res")
             args = []
             for arg in expr.args:
                ar = self.genExprCode(arg)
                args.append(ar)
             fn = self.funcMap[expr.proc]
             ins = ir.Call(fn, args, tmp)
             self.builder.addIns(ins)
             return tmp
        else:
             raise NotImplementedError('Unknown expr {}'.format(expr))