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


class CodeGenerator(ir.Builder):
    """ 
      Generates intermediate (IR) code from a package. The entry function is
      'genModule'. The main task of this part is to rewrite complex control
      structures, such as while and for loops into simple conditional
      jump statements. Also complex conditional statements are simplified.
      Such as 'and' and 'or' statements are rewritten in conditional jumps.
      And structured datatypes are rewritten.
    """
    def __init__(self):
        self.logger = logging.getLogger('c3cgen')

    def gencode(self, pkg):
        self.prepare()
        assert type(pkg) is astnodes.Package
        self.logger.info('Generating ir-code for {}'.format(pkg.name))
        self.varMap = {} # Maps variables to storage locations.
        self.funcMap = {}
        self.m = ir.Module(pkg.name)
        self.genModule(pkg)
        return self.m

    # inner helpers:
    def genModule(self, pkg):
        # Take care of forward declarations:
        for s in pkg.innerScope.Functions:
            f = self.newFunction(s.name)
            self.funcMap[s] = f
        for v in pkg.innerScope.Variables:
            self.varMap[v] = self.newTemp()
        for s in pkg.innerScope.Functions:
            self.genFunction(s)

    def genFunction(self, fn):
        # TODO: handle arguments
        f = self.funcMap[fn]
        f.return_value = self.newTemp()
        # TODO reserve room for stack, this can be done at later point?
        self.setFunction(f)
        l2 = self.newBlock()
        self.emit(ir.Jump(l2))
        self.setBlock(l2)
        # generate room for locals:

        for sym in fn.innerScope:
            # TODO: handle parameters different
            if sym.isParameter:
                v = ir.Parameter(sym.name)
                f.addParameter(v)
            elif sym.isLocal:
                v = ir.LocalVariable(sym.name)
                f.addLocal(v)
            else:
                #v = self.newTemp()
                raise NotImplementedError('{}'.format(sym))
            # TODO: make this ssa here??
            self.varMap[sym] = v

        self.genCode(fn.body)
        # Set the default return value to zero:
        # TBD: this may not be required?
        self.emit(ir.Move(f.return_value, ir.Const(0)))
        self.emit(ir.Jump(f.epiloog))
        self.setFunction(None)

    def genCode(self, code):
        assert isinstance(code, astnodes.Statement)
        self.setLoc(code.loc)
        if type(code) is astnodes.CompoundStatement:
            for s in code.statements:
                self.genCode(s)
        elif type(code) is astnodes.Assignment:
            rval = self.genExprCode(code.rval)
            lval = self.genExprCode(code.lval)
            self.emit(ir.Move(lval, rval))
        elif type(code) is astnodes.ExpressionStatement:
            self.emit(ir.Exp(self.genExprCode(code.ex)))
        elif type(code) is astnodes.IfStatement:
            bbtrue = self.newBlock()
            bbfalse = self.newBlock()
            te = self.newBlock()
            self.genCondCode(code.condition, bbtrue, bbfalse)
            self.setBlock(bbtrue)
            self.genCode(code.truestatement)
            self.emit(ir.Jump(te))
            self.setBlock(bbfalse)
            if code.falsestatement:
                self.genCode(code.falsestatement)
            self.emit(ir.Jump(te))
            self.setBlock(te)
        elif type(code) is astnodes.ReturnStatement:
            if code.expr:
                re = self.genExprCode(code.expr)
                self.emit(ir.Move(self.fn.return_value, re))
                self.emit(ir.Jump(self.fn.epiloog))
                b = self.newBlock()
                self.setBlock(b)
            else:
                self.builder.addIns(ir.Return())
        elif type(code) is astnodes.WhileStatement:
            bbdo = self.newBlock()
            bbtest = self.newBlock()
            te = self.newBlock()
            self.emit(ir.Jump(bbtest))
            self.setBlock(bbtest)
            self.genCondCode(code.condition, bbdo, te)
            self.setBlock(bbdo)
            self.genCode(code.statement)
            self.emit(ir.Jump(bbtest))
            self.setBlock(te)
        else:
            raise NotImplementedError('Unknown stmt {}'.format(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.newBlock()
                self.genCondCode(expr.a, bbtrue, l2)
                self.setBlock(l2)
                self.genCondCode(expr.b, bbtrue, bbfalse)
            elif expr.op == 'and':
                l2 = self.newBlock()
                self.genCondCode(expr.a, l2, bbfalse)
                self.setBlock(l2)
                self.genCondCode(expr.b, bbtrue, bbfalse)
            elif expr.op in ['==', '>', '<']:
                ta = self.genExprCode(expr.a)
                tb = self.genExprCode(expr.b)
                self.emit(ir.CJump(ta, expr.op, tb, bbtrue, bbfalse))
            else:
                raise NotImplementedError('Unknown condition {}'.format(expr))
        elif type(expr) is astnodes.Literal:
             if expr.val:
                self.emit(ir.Jump(bbtrue))
             else:
                self.emit(ir.Jump(bbfalse))
        else:
             raise NotImplementedError('Unknown cond {}'.format(expr))

    def genExprCode(self, expr):
        assert isinstance(expr, astnodes.Expression)
        if type(expr) is astnodes.Binop and expr.op in ir.Binop.ops:
            ra = self.genExprCode(expr.a)
            rb = self.genExprCode(expr.b)
            return ir.Binop(ra, expr.op, rb)
        elif type(expr) is astnodes.Unop and expr.op == '&':
            ra = self.genExprCode(expr.a)
            assert type(ra) is ir.Mem
            return ra.e
        elif type(expr) is astnodes.VariableUse:
            # This returns the dereferenced variable.
            return ir.Mem(self.varMap[expr.target])
        elif type(expr) is astnodes.Deref:
            # dereference pointer type:
            addr = self.genExprCode(expr.ptr)
            return ir.Mem(addr)
        elif type(expr) is astnodes.FieldRef:
            base = self.genExprCode(expr.base)
            assert type(base) is ir.Mem, type(base)
            base = base.e
            bt = theType(expr.base.typ)
            offset = ir.Const(bt.fieldOffset(expr.field))
            return ir.Mem(ir.Add(base, offset))
        elif type(expr) is astnodes.Literal:
            return ir.Const(expr.val)
        elif type(expr) is astnodes.TypeCast:
            # TODO: improve this mess:
            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:
            args = [self.genExprCode(e) for e in expr.args]
            fn = self.funcMap[expr.proc]
            return ir.Call(fn, args)
        else:
            raise NotImplementedError('Unknown expr {}'.format(expr))