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 code from a package """
    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:
            #print(v)
            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:
                print('param', sym)
                ir.Parameter(sym.name)
            if sym.isLocal:
                print('local', sym)
            
            v = self.newTemp()
            # 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.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))
            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)
            # TODO: Make this work?
            return ra
        elif type(expr) is astnodes.VariableUse:
            return 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:
            b = self.genExprCode(expr.base)
            #b = b.e
            bt = theType(expr.base.typ)
            offset = ir.Const(bt.fieldOffset(expr.field))
            return ir.Mem(ir.Binop(b, '+', 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))