from back.tac import *
from back.generator import Register
import sys

class Node(object):
    labels = 0

    def newlabel(self):
        Node.labels += 1
        return Node.labels

    def emitlabel(self, i):
        TACList().add(Label("L%d" % i))

    def emit(self, op, arg1=None, arg2=None):
        self.debug("emit(): op = %s, arg = %s, arg2 = %s" % (op, arg1, arg2))
        TACList().add(TAC(op, arg1, arg2))

    def debug(self, msg):
        #sys.stdout.write("DEBUG: " + msg + "\n")
        #sys.stdout.flush()
        return

class Program(Node):
    def __init__(self, functions, lineno = -1):
        self.functions = functions
        self.lineno = lineno

    def generate(self):
        self.debug("Program.generate(): %s" % repr(self))
        self.emit(Op.JMP, "main")
        for function in self.functions:
            function.generate()

    def __repr__(self):
        return "<Program: %s>" % self.functions

class Function(Node):
    def __init__(self, name, params, statement, lineno = -1):
        self.name = name
        self.params = params
        self.statement = statement
        self.lineno = lineno

    def emit_prologue(self, name):
        TACList().add(FunctionPrologue(name))

    def emit_epilogue(self, name):
        TACList().add(FunctionEpilogue(name))

    def generate(self):
        self.debug("Function.generate(): %s" % repr(self))
        self.emit_prologue(self.name)
        Register().set_function(self.name)
        begin = self.newlabel()
        after = self.newlabel()
        self.emitlabel(begin)
        self.statement.generate(begin, after, after)
        self.emit(Op.RETURN, 0, after)
        self.emitlabel(after)
        self.emit_epilogue(self.name)

    def __repr__(self):
        return "<Function: %s(%s) at line %d: %s>" % (self.name, str(self.params), self.lineno, str(self.statement))

class Statement(Node):
    pass

class Expression(Node):
    pass

class Sequence(Statement):
    def __init__(self, s1, s2, lineno = -1):
        self.s1 = s1
        self.s2 = s2
        self.lineno = lineno

    def generate(self, before, after, last):
        self.debug("Sequence.generate(before = %d, after = %d): %s" % (before, after, repr(self)))
        if not self.s1:
            self.s2.generate(before, after, last)
        elif not self.s2:
            self.s1.generate(before, after, last)
        else:
            label = self.newlabel()
            self.s1.generate(before, label, last)
            self.emitlabel(label)
            self.s2.generate(label, after, last)

    def __repr__(self):
        return "<Sequence at line %d: %s, %s>" % (self.lineno, self.s1, self.s2)

class IfStatement(Statement):
    def __init__(self, expression, true_statement, false_statement, lineno = -1):
        self.expression = expression
        self.true_statement = true_statement
        self.false_statement = false_statement
        self.lineno = lineno

    def generate(self, before, after, last):
        self.debug("IfStatement.generate(before = %d, after = %d): %s" % (before, after, repr(self)))

        x = self.expression.reduce()

        if self.false_statement:
            true_label = self.newlabel()
            false_label = self.newlabel()
            self.emit(Op.BEZ, x, "L%d" % false_label)
            self.emitlabel(true_label)
            self.true_statement.generate(true_label, after, last)
            self.emit(Op.JMP, "L%d" % after)
            self.emitlabel(false_label)
            self.false_statement.generate(false_label, after, last)
        else:
            true_label = self.newlabel()
            self.emit(Op.BEZ, x, "L%d" % after)
            self.emitlabel(true_label)
            self.true_statement.generate(true_label, after, last)

    def __repr__(self):
        return "<IfStatement at line %d>" % self.lineno

class WhileStatement(Statement):
    def __init__(self, expression, statement, lineno = -1):
        self.expression = expression
        self.statement = statement
        self.lineno = lineno

    def generate(self, before, after, last):
        self.debug("WhileStatement.generate(before = %d, after = %d): %s" % (before, after, repr(self)))
        self.emit(Op.BEZ, self.expression.reduce(), "L%d" % after)
        label = self.newlabel()
        self.emitlabel(label)
        self.statement.generate(label, before, last)
        self.emit(Op.JMP, "L%d" % before)

    def __repr__(self):
        return "<WhileStatement at line %d>" % self.lineno

class ReturnStatement(Statement):
    def __init__(self, expression, lineno = -1):
        self.expression = expression
        self.lineno = lineno

    def generate(self, before, after, last):
        self.debug("ReturnStatement.generate(before = %d, after = %d): %s" % (before, after, repr(self)))
        self.emit(Op.RETURN, self.expression.reduce(), last)

    def __repr__(self):
        return "<ReturnStatement at line %d>" % self.lineno

class AssignStatement(Statement):
    def __init__(self, ident, expression, lineno = -1):
        self.ident = ident
        self.expression = expression
        self.lineno = lineno

    def generate(self, before, after, last):
        self.debug("AssignStatement.generate(before = %d, after = %d): %s" % (before, after, repr(self)))
        self.emit(Op.STORE, self.ident, self.expression.reduce())

    def __repr__(self):
        return "<AssignStatement for %s at line %d>" % (self.ident, self.lineno)

class FunctionCall(Statement, Expression):
    def __init__(self, ident, arguments = [], lineno = -1):
        self.ident = ident
        self.arguments = arguments
        self.lineno = lineno

    def push_params(self):
        self.arguments.reverse()
        for argument in self.arguments:
            self.emit(Op.PUSH, argument.reduce())
        self.arguments.reverse()

    def pop_params(self):
        for argument in self.arguments:
            self.emit(Op.POP, 0)

    # XXX: wrong signature?
    def generate(self):
        self.debug("FunctionCall.generate(): %s" % (repr(self)))
        self.push_params()
        self.emit(Op.CALL, self.ident, 0)
        self.pop_params()

    def reduce(self):
        self.debug("FunctionCall.reduce(): %s" % repr(self))
        self.push_params()
        r = Register().new()
        self.emit(Op.CALL, self.ident, r)
        self.pop_params()
        return r

    def __repr__(self):
        return "<FunctionCall for %s(%s) at line %d>" % (self.ident, self.arguments, self.lineno)

class Operation(Expression):
    def __init__(self, op, lineno = -1):
        self.op = op
        self.lineno = -1

    def __repr__(self):
        return "<Operation at line %d>" % self.lineno

class UnaryExpression(Operation):
    """ Unary expressions are !x or -x ..."""
    def __init__(self, op, right, lineno = -1):
        super(UnaryExpression, self).__init__(op, lineno)
        self.right = right

    def reduce(self):
        self.debug("UnaryExpression.generate(): %s" % repr(self))
        op = {
                "!": Op.NOT,
                "-": Op.MINUS,
             }[self.op]
        r = self.right.reduce()
        self.emit(op, r)
        return r

    def __repr__(self):
        return "<UnaryExpression: %s(%s) at line %d>" % (self.op, self.right, self.lineno)

class BinaryExpression(Operation):
    """ Biary expression are x+y, x-y, ... """
    def __init__(self, left, op, right, lineno = -1):
        super(BinaryExpression, self).__init__(op, lineno)
        self.left = left
        self.right = right

    def reduce(self):
        self.debug("BinaryExpression.reduce(): %s" % repr(self))
        arith_op = {
                "+": Op.ADD,
                "-": Op.SUB,
                "*": Op.MUL,
                "/": Op.DIV,
                "%": Op.MOD,
             }

        logic_op = {
                "==": Op.EQ,
                "!=": Op.NE,
                "<":  Op.LT,
                "<=": Op.LE,
                ">=": Op.GE,
                ">":  Op.GT,
             }

        try:
            op = arith_op[self.op]
            r = self.left.reduce()
            self.emit(op, r, self.right.reduce())
        except KeyError:
            op = logic_op[self.op]
            r = Register().new()
            self.emit(Op.CMP, self.left.reduce(), self.right.reduce())
            self.emit(op, r)

        return r

    def __repr__(self):
        return "<BinaryExpression: %s(%s, %s) at line %d>" % (self.op, self.left, self.right, self.lineno)

class Variable(Expression):
    def __init__(self, name, lineno = -1):
        self.name = name
        self.lineno = lineno

    def reduce(self):
        self.debug("Variable.reduce(): %s" % repr(self))
        r = Register().new()
        self.emit(Op.LOAD, self.name, r)
        return r

    def __repr__(self):
        return "<Variable: %s at line %d>" %(self.name, self.lineno)

class Constant(Expression):
    def __init__(self, value, lineno = -1):
        self.value = value
        self.lineno = lineno

    def reduce(self):
        self.debug("Constant.reduce(): %s" % repr(self))
        r = Register().new()
        self.emit(Op.MOV, r, self.value)
        return r

    def __repr__(self):
        return "<Constant: %d at line %d>" %(self.value, self.lineno)