import io, os.path, re

from . import cparse
from .cvisitors import Visitor

##
## STACK MACHINE ABSTRACTION
##

class TtcRegisters (object):
    def __init__(self, parent, base_fp):
        # we keep R8 and R9 for immediate temp storage
        self.all_regs = ['R%s' % i for i in range(8)]
        self.regs_free = self.all_regs[:]
        self.regs_almost_free = []
        self.mem_free = []
        self.stack = []
        self.next_temp = base_fp - 1
        self.parent = parent
        self.callee_save_regs_used = []
        self.caller_save_regs = ['R%s' % i for i in range(4)]
        self.callee_save_regs = ['R%s' % i for i in range(4, 8)]

    def o(self, str, comment=None):
        self.parent.o(str, comment)

    def save_caller_saves(self):
        for reg in self.caller_save_regs:
            if reg not in self.regs_free:
                self._copy_reg_to_temp([reg], "Save caller-save register to temp")
                self.regs_free.append(reg)

    def save_callee_saves(self):
        for reg in self.callee_save_regs_used:
            self.o("  push %s" % reg, "Save callee-save register")

    def load_callee_saves(self):
        for reg in reversed(self.callee_save_regs_used):
            self.o("  pop  %s" % reg, "Restore callee-save register")

    def _copy_reg_to_temp(self, valid_regs, comment_str=None):
        if len(self.mem_free) == 0:
            self.mem_free.append(self.next_temp)
            self.next_temp -= 1
        mem = self.mem_free.pop()
        reg = None
        index = 0
        for i in self.stack:
            if i in valid_regs:
                reg = i
                break
            index += 1
        if reg == None:
            raise Exception("No free registers inside OR outside of stack!")
        if comment_str == None:
            comment_str = "Stack machine: copy register to temp"
        self.o("  sti %s %s" % (reg, mem), comment_str)
        self.stack[index] = mem
        return reg

    def _get_free_reg(self, valid_regs, preferred_reg=None):
        if len(self.regs_free) > 0:
            reg = None
            if preferred_reg != None and preferred_reg in self.regs_free:
                reg = preferred_reg
            else:
                for r in self.regs_free:
                    if r in valid_regs:
                        reg = r
            if reg != None:
                self.regs_free.remove(reg)
                if reg in self.callee_save_regs and reg not in self.callee_save_regs_used:
                    self.callee_save_regs_used.append(reg)
                return reg
        return self._copy_reg_to_temp(valid_regs)

    def push(self, preferred_reg=None, valid_regs=None):
        if valid_regs == None:
            valid_regs = self.all_regs
        reg = self._get_free_reg(valid_regs, preferred_reg)
        self.stack.append(reg)
        return reg

    def pop(self, valid_regs=None):
        if valid_regs == None:
            valid_regs = self.all_regs
        return self._pop(valid_regs)

    def _pop(self, valid_regs):
        loc = self.stack.pop()
        if loc in valid_regs:
            self.regs_almost_free.append(loc)
            return loc
        reg = self._get_free_reg(valid_regs)
        self.o("  ldi %s %s" % (reg, loc),
               "Stack machine: copy temp to register")
        if loc in self.all_regs:
            self.regs_free.append(loc)
        self.regs_almost_free.append(reg)
        return reg

    def peek(self):
        return self.stack[-1]

    def is_empty(self):
        return len(self.stack) == 0

    def done(self):
        self.regs_free.extend(self.regs_almost_free)
        self.regs_almost_free = []

    def get_max_fp(self):
        return self.next_temp + 1

##
## CODE GENERATOR
##

class CodeGenVisitor (Visitor) :
    __label = 0
    __str_literal_label = 0

    def __init__(self, path, test=None):
        Visitor.__init__(self)
        self.path = path
        self._lbl = re.sub("[^a-z0-9]", "", os.path.splitext(path)[0],
                           flags=re.I)
        self.test = test
        self.str_literal_str = io.StringIO()
        self.str_literal_dict = {}

    def new_label(self):
        label = "_%s_L%d" % (self._lbl, self.__label)
        self.__class__.__label += 1
        return label

    def o(self, str, comment=None):
        if comment is not None:
            comment = "# %s" % comment
            self.curr_str.write("%-35s %s\n" % (str, comment))
        elif str :
            self.curr_str.write(str + "\n")

    def c(self, str, indent_amt=2):
        self.curr_str.write("\n  # %s\n\n" % str)

    def _loc(self, node, comment="") :
        if node.lineno != self._last_lineno or comment :
            self._last_lineno = node.lineno
            self.curr_str.write("# %s %s\n" % (node.loc(), comment))

    def vNodeList(self, node):
        self._visitList(node.nodes)

    def _empty_stack(self, node):
        if not self.stack.is_empty():
            self.stack.pop()
            self.stack.done()
            if not self.stack.is_empty():
                raise Exception("PANIC! Register stack isn't empty!")

    def _accept_and_empty_stack(self, node):
        node.accept(self)
        self._empty_stack(node)

    def vStatementList(self, node):
        for n in node.nodes:
            self._loc(n)
            self._accept_and_empty_stack(n)

    def _generate_global_variable_definitions(self, node):
        globals_str = io.StringIO()
        for symbol in node.symtab.entries.values():
            symbol.compile_loc = symbol.name
            if not symbol.type.is_function() and not symbol.extern:
                globals_str.write("%s:\n  int 0\n" % symbol.compile_loc)
        return globals_str

    def vTranslationUnit(self, node):
        self._last_lineno = None
        self.curr_str = io.StringIO()
        self.globals_str = self._generate_global_variable_definitions(node)
        self._visitList(node.nodes)

    @classmethod
    def concat (cls, boot, outfile, chunks) :
        outfile.write("# Generated by cct\n"
                      "# Franck Pommereau (2018)\n"
                      "# Adapted from Atul Varma's c.py (Spring 2004)\n\n")
        if boot :
            outfile.write("#\n"
                          "# on computer start: call main and halt\n"
                          "#\n\n"
                          "IRQ0:\n"
                          "  set  R9 @%s\n"
                          "  call R9\n"
                          "  halt\n\n" % boot)
        for c in chunks :
            s = c.curr_str.getvalue()
            if s.strip() :
                outfile.write("#\n# code from file %r\n#\n" % c.path)
                outfile.write(s)
        for c in chunks :
            s = c.globals_str.getvalue()
            if s.strip() :
                outfile.write("\n#\n# globals from file %r\n#\n\n" % c.path)
                outfile.write(s)
        for c in chunks :
            s = c.str_literal_str.getvalue()
            if s.strip() :
                outfile.write("\n#\n# string literals from file %r\n#\n\n"
                              % c.path)
                outfile.write(s)

    def _calc_function_var_addrs(self, symtab, last_fp_loc):
        self._calc_function_arg_addrs(symtab)
        return self._calc_local_var_addrs(symtab.children[0], last_fp_loc)

    def _calc_function_arg_addrs(self, symtab):
        for symbol in symtab.entries.values():
            symbol.compile_loc = 3 + symbol.param_num
            if not symbol.is_used:
                self.warning("function argument '%s' is never used." % symbol.name)

    def _calc_local_var_addrs(self, symtab, last_fp_loc):
        for symbol in symtab.entries.values():
            if symbol.extern:
                symbol.compile_loc = "@" + symbol.name
                continue
            last_fp_loc -= 1
            symbol.compile_loc = last_fp_loc
            if not symbol.is_used:
                self.warning("local variable '%s' is never used." % symbol.name)
        max_last_fp = last_fp_loc
        for kid in symtab.children:
            curr_last_fp = self._calc_local_var_addrs(kid, last_fp_loc)
            if curr_last_fp < max_last_fp:
                max_last_fp = curr_last_fp
        return max_last_fp

    def _fill_line(self, str, width=70):
        extra = "-" * (width-1-len(str))
        return str + " " + extra

    def vFunctionDefn(self, node):
        self.break_labels = []
        self.continue_labels = []
        self.curr_func_end_label = self.new_label() + "_function_end"
        stack_frame_size = self._calc_function_var_addrs(node.symtab, 0)
        line = self._fill_line("BEGIN FUNCTION: %s()" % node.name)
        self.c("%s\n  #\n  # Function type: %s" % (line, node.type.get_string()), 0)
        self._loc(node, "(enter function)")
        self.o("%s:" % node.compile_loc)
        self.o("  push BP", "Save old frame pointer")
        self.o("  mov  SP BP", "Set new frame pointer")
        self.stack = TtcRegisters(self, stack_frame_size)
        old_str = self.curr_str
        self.curr_str = io.StringIO()
        node.body.accept(self)
        function_str = self.curr_str
        self.curr_str = old_str
        frame_size = abs(self.stack.get_max_fp())
        if frame_size != 0:
            self.o("  set  R9 %X" % frame_size,
                   "Allocate %s words for local+temp vars" % frame_size)
            self.o("  sub  SP R9 SP", "... shift SP")
        self.stack.save_callee_saves()
        self.curr_str.write(function_str.getvalue())
        self._loc(node, "(exit function)")
        self.o("%s:" % self.curr_func_end_label)
        self.stack.load_callee_saves()
        self.o("  mov  BP SP", "Deallocate local+temp vars")
        self.o("  pop  BP", "Restore old stack frame")
        self.o("  ret")
        line = self._fill_line("END FUNCTION: %s()" % node.name)
        self.c(line, 0)

    def vCompoundStatement(self, node):
        self._loc(node)
        node.statement_list.accept(self)

    def vIfStatement(self, node):
        self._loc(node)
        done_label = self.new_label() + "_done"
        if not node.else_stmt.is_null():
            else_label = self.new_label() + "_else"
        else:
            else_label = done_label
        self.c("IF statment - begin")
        node.expr.accept(self)
        comparer = self.stack.pop()
        self.stack.done()
        self.o("  set  R9 @%s" % else_label, "Point towards else clause")
        self.o("  jz   %s R9" % comparer, "... if result is zero, jump to it")
        self.c("IF statment - THEN clause - begin")
        self._accept_and_empty_stack(node.then_stmt)
        self.c("IF statment - THEN clause - end")
        self.o("  set  R9 @%s" % done_label, "Point towards if end")
        self.o("  jmp  R9", "... jump to it")
        self.c("IF statment - ELSE clause - begin")
        self.o("%s:" % else_label)
        if not node.else_stmt.is_null():
            self._accept_and_empty_stack(node.else_stmt)
            self.c("IF statment - ELSE clause - end")
        self.o("%s:" % done_label)
        self.c("IF statment - end")

    def _push_loop_labels(self, break_label, continue_label):
        self.break_labels.append(break_label)
        self.continue_labels.append(continue_label)

    def _pop_loop_labels(self):
        self.break_labels.pop()
        self.continue_labels.pop()

    def vWhileLoop(self, node):
        self._loc(node)
        test_label = self.new_label() + "_test"
        done_label = self.new_label() + "_done"
        self._push_loop_labels(break_label=done_label,
                               continue_label=test_label)
        self.c("WHILE loop - begin")
        self.o("%s:" % test_label)
        node.expr.accept(self)
        comparer = self.stack.pop()
        self.stack.done()
        self.o("  set  R9 @%s" % done_label, "Point towards loop exit")
        self.o("  jz   %s R9" % comparer, "... if result is zero, jump to it")
        self._accept_and_empty_stack(node.stmt)
        self.o("  set  R9 @%s" % test_label, "Point towards loop start")
        self.o("  jmp  R9", "... jump to it")
        self.o("%s:" % done_label)
        self.c("WHILE loop - end")
        self._pop_loop_labels()

    def vForLoop(self, node):
        self._loc(node)
        test_label = self.new_label() + "_test"
        done_label = self.new_label() + "_done"
        self._push_loop_labels(break_label=done_label,
                               continue_label=test_label)
        self.c("FOR loop - begin")
        self._accept_and_empty_stack(node.begin_stmt)
        self.o("%s:" % test_label)
        node.expr.accept(self)
        comparer = self.stack.pop()
        self.stack.done()
        self.o("  set  R9 @%s" % done_label, "Point towards loop exit")
        self.o("  jz   %s R9" % comparer, "... if result is zero, jump to it")
        self._accept_and_empty_stack(node.stmt)
        self._accept_and_empty_stack(node.end_stmt)
        self.o("  set  R9 @%s" % test_label, "Point towards while loop start")
        self.o("  jmp  R9", "... jump to it")
        self.o("%s:" % done_label)
        self.c("FOR loop - end")
        self._pop_loop_labels()

    def vBreakStatement(self, node):
        self._loc(node)
        self.o("  set  R9 @%s" % self.break_labels[-1], "Loop: break statement")
        self.o("  jmp  R9", "... jump to loop exit")

    def vContinueStatement(self, node):
        self._loc(node)
        self.o("  set  R9 @%s" % self.continue_labels[-1], "Loop: continue statement")
        self.o("  jmp  R9", "... jump to loop start")

    def _get_new_str_literal_label(self, str):
        if str in self.str_literal_dict :
            return self.str_literal_dict[str]

        label_str = "_%s_LC%d" % (self._lbl, self.__str_literal_label)
        self.str_literal_dict[str] = label_str
        str = str.replace('\n', '\\n')
        self.str_literal_str.write("""%s:\n  str "%s\\0"\n""" % (label_str, str))
        self.__class__.__str_literal_label += 1
        return label_str

    def vStringLiteral(self, node):
        self._loc(node)
        label_str = self._get_new_str_literal_label(node.get_str())
        COMMENT_CHARS = 7
        comment_label = node.get_sanitized_str()
        if len(comment_label) > COMMENT_CHARS:
            comment_label = "%s..." % comment_label[0:COMMENT_CHARS]
        self.o("  set  %s @%s" % (self.stack.push(), label_str),
               "Get addr of string literal '%s'" % comment_label)

    def vConst(self, node):
        self._loc(node)
        self.o("  set  %s %X" % (self.stack.push(), node.value),
               "Load numeric constant %d" % node.value)

    def vId(self, node):
        self._loc(node)
        if node.output_addr:
            if isinstance(node.symbol.compile_loc, int) :
                self.o("  set  R9 %X" % node.symbol.compile_loc,
                       "Get BP-relative address of %s" % node.symbol.name)
                self.o("  add  R9 BP %s" % self.stack.push(),
                       "Compute address of %s" % node.symbol.name)
            else :
                self.o("  set  %s @%s" % (self.stack.push(), node.symbol.compile_loc),
                       "Get address of %s" % node.symbol.name)
            return
        reg = self.stack.push()
        if isinstance(node.symbol.compile_loc, int) :
            self.o("  ldi  %s %s" % (reg, node.symbol.compile_loc),
                   "Get value of %s" % node.symbol.name)
        else :
            self.o("  set  %s @%s" % (reg, node.symbol.compile_loc),
                   "Get address of %s" % node.symbol.name)
            self.o("  ld   %s %s" % (reg, reg), "Get value of %s" % node.symbol.name)

    def vArrayExpression(self, node):
        self._loc(node)
        node.expr.accept(self)
        node.index.accept(self)
        reg_index = self.stack.pop()
        reg_expr = self.stack.pop()
        reg_to = self.stack.push()
        self.stack.done()
        self.o("  add  %s %s %s" % (reg_expr, reg_index, reg_to),
               "Load addr of array index")
        if not node.output_addr:
            self.o("  ld   %s %s" % (reg_to, reg_to), "Load array value")

    def vFunctionExpression(self, node):
        self._loc(node)
        if self.test and node.function.symbol.name == "TTCTEST" :
            call = node.arglist.nodes[0].get_str()
            expect = self.test[call]
            ret = self._accept_and_pop(node.arglist.nodes[1])
            self.o("  ### %s == %s # == %s\n" % (ret, expect, call))
            self.stack.done()
            return
        self.c("FUNCTION CALL to %s() - begin" %
               node.function.symbol.name)
        self.stack.save_caller_saves()
        argnum = len(node.arglist.nodes)
        for arg in reversed(node.arglist.nodes):
            arg_reg = self._accept_and_pop(arg)
            self.o("  push %s" % arg_reg, "Push arg %d" % argnum)
            self.stack.done()
            argnum -= 1
        self.o("  set  R9 @%s" % node.function.symbol.compile_loc,
               "Point towards function %s()" % node.function.symbol.compile_loc)
        self.o("  call R9", "... call it")
        result = self.stack.push(preferred_reg='R0')
        if result != 'R0':
            self.o("  mov  R0 %s" % result, "Copy return value")
        arg_stack_size = len(node.arglist.nodes)
        if arg_stack_size > 0:
            self.o("  set  R9 %X" % arg_stack_size, "Deallocate argument stack")
            self.o("  add  R9 SP SP", "... shift SP")
        self.c("FUNCTION CALL to %s() - end" % node.function.symbol.name)

    def vReturnStatement(self, node):
        self._loc(node)
        return_reg = self._accept_and_pop(node.expr)
        self.o("  mov  %s R0" % return_reg, "Set return value")
        self.o("  set  R9 @%s" % self.curr_func_end_label,
               "Point towards function exit")
        self.o("  jmp  R9", "... jump to it")
        self.stack.done()

    def _accept_and_pop(self, node):
        if node.is_const():
            self.o("  set  R8 %X" % node.value, "Use constant %s" % node.value)
            return "R8"
        else:
            node.accept(self)
            return self.stack.pop()

    def _binop_assign(self, node):
        node.left.accept(self)
        right_reg = self._accept_and_pop(node.right)
        left_reg = self.stack.pop()
        reg = self.stack.push()
        if node.op == "=" :
            self.o("  mov  %s %s" % (right_reg, reg),
                   "Assignement '=': set result")
        elif node.op == "+=" :
            self.o("  ld   %s %s" % (left_reg, reg),
                   "Assignement '+=': load initial left value")
            self.o("  add  %s %s %s" % (reg, right_reg, reg),
                   "... add right value")
        elif node.op == "-=" :
            self.o("  ld   %s %s" % (left_reg, reg),
                   "Assignement '-=': load initial left value")
            self.o("  sub  %s %s %s" % (reg, right_reg, reg),
                   "... subtract right value")
        else :
            raise Exception("unsupported assignment %r" % node.op)
        self.o("  st   %s %s" % (reg, left_reg),
               "... save value")
        self.stack.done()

    _binary_operations = {"+" : "add",
                          "-" : "sub",
                          "*" : "mul",
                          "/" : "div",
                          "%" : "mod",
                          "<<" : "lshift",
                          ">>" : "rshift",
                          "&&" : "and",
                          "||" : "or",
                          "&" : "and",
                          "|" : "or",
                          "^" : "xor"}
    def _binop_arith(self, node):
        node.left.accept(self)
        right_reg = self._accept_and_pop(node.right)
        left_reg = self.stack.pop()

        try :
            instr = self._binary_operations[node.op]
        except KeyError :
            raise Exception("unsupported operation %r" % node.op)

        self.o("  %s %s %s %s" % (instr.ljust(4), left_reg, right_reg, left_reg),
               "Perform '%s'" % node.op)
        self.stack.done()
        new_reg = self.stack.push(preferred_reg=left_reg)
        if new_reg != left_reg:
            raise Exception("PANIC! Binop push() isn't same as last pop()!")

    def _binop_compare(self, node):
        node.left.accept(self)
        right_reg = self._accept_and_pop(node.right)
        left_reg = self.stack.pop()
        self.stack.done()
        reg = self.stack.push()
        if node.op == "==" :
            self.o("  eq   %s %s %s" % (left_reg, right_reg, reg), "Perform '=='")
        elif node.op == "!=" :
            self.o("  eq   %s %s %s" % (left_reg, right_reg, reg), "Perform '!='")
            self.o("  not  %s %s" % (reg, reg), "... throught 'not =='")
        elif node.op == ">" :
            self.o("  gt   %s %s %s" % (left_reg, right_reg, reg), "Perform '>'")
        elif node.op == "<" :
            self.o("  gt   %s %s %s" % (right_reg, left_reg, reg), "Perform '<'")
        elif node.op == ">=" :
            self.o("  inc  %s" % left_reg, "Perform 'a >= b'")
            self.o("  gt   %s %s %s" % (left_reg, right_reg, reg),
                   "... that is 'a+1 > b'")
        elif node.op == "<=" :
            self.o("  inc  %s" % right_reg, "Perform 'a <= b'")
            self.o("  gt   %s %s %s" % (right_reg, left_reg, reg),
                   "... that is 'b+1 > a'")
        else :
            raise Exception("unsupported comparison %r" % node.op)

    def vBinop(self, node):
        self._loc(node)
        if node.op in cparse.Binop.ASSIGN_OPS:
            self._binop_assign(node)
        elif node.op in ['+','-','*','/','%']:
            self._binop_arith(node)
        elif node.op in ['==', '!=', '<', '>', '<=', '>=']:
            self._binop_compare(node)
        else :
            raise Exception("unsupported operation %r" % node.op)

    def vNegative(self, node):
        self._loc(node)
        node.expr.accept(self)
        reg = self.stack.peek()
        self.o("  neg  %s %s" % (reg, reg), "Perform unary negation")

    def vPointer(self, node):
        self._loc(node)
        node.expr.accept(self)
        if node.output_addr:
            self.o("", "(Getting pointer target addr via '*')")
            return
        reg_from = self.stack.pop()
        reg_to = self.stack.push()
        self.o("  ld   %s %s" % (reg_from, reg_to), "Pointer dereference")
        self.stack.done()

    def vAddrOf(self, node):
        self._loc(node)
        node.expr.accept(self)
        self.o("", "(Address-of operator '&' used here)")

def compile (path) :
    tgt = os.path.splitext(path)[0] + ".bios"
    tgt, ["ttc", "asm", path, tgt]

def cpp (args) :
    if args.test :
        return ["-DTTC"]
    else :
        return []