import tools.wrapper as w
import argparse

class CallState:
    def __init__(self, machine, addr):
        self.ra = machine.getRegister(1)
        self.sp = machine.getRegister(2)
        self.saved_registers = [machine.getRegister(i) for i in [8, 9, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27]]
        self.pc = machine.getPC
        self.addr = addr # address of the called function or the return address

def check_cc(call_state, return_state):
    """Given the state of registers when a function was called and when the
    function is returned from, validate the RISC-V calling convention and
    return a list of lines to print to report errors."""

    error_lines = []

    if call_state.sp != return_state.sp:
        if call_state.sp < return_state.sp:
            error_lines.append("WARNING: The stack pointer increased (moved downwards) during a call to 0x{:x}".format(call_state.addr))
        else:
            error_lines.append("WARNING: The stack pointer decreased (moved upwards) during a call to 0x{:x}".format(call_state.addr))
        error_lines.append("  Check that you correctly decreased and restored the sp register.")
    if return_state.ra != call_state.pc + 4:
        error_lines.append("WARNING: The code is returning to an address (0x{:x}) different from the address of the instruction after it was called (0x{:x})".format(return_state.ra, call_state.pc + 4))
        error_lines.append("  Check that you are correctly storing and restoring the ra register.")
    if call_state.saved_registers != return_state.saved_registers:
        changes = [(i, v1, v2) for i, (v1, v2) in enumerate(zip(call_state.saved_registers, return_state.saved_registers)) if v1 != v2]
        changed_registers = ", ".join("s{}".format(i) for i, _1, _2 in changes)
        if len(changes) == 1:
            error_lines.append("WARNING: The value of saved register {} was not preserved during a call to 0x{:x}".format(changed_registers, call_state.addr))
        else:
            error_lines.append("WARNING: The values of saved registers {} were not preserved during a call to 0x{:x}".format(changed_registers, call_state.addr))
        for i, v1, v2 in changes:
            error_lines.append("    s{}: {} -> {}".format(i, v1, v2))
        error_lines.append("  Check that you are correctly storing and restoring saved registers.")

    return error_lines

# function to monkey patch with
import tools.Riscv as Riscv
def execute_wrapper(orig_execute):
    call_states = []
    total_error_lines = [0] # hack for mutability

    def execute(self):
        error_lines = []
        if self.opcode == Riscv.opcode_STORE:
            addr = self.getRegister(self.rs1) + self.simm12
            if addr < 0x100: # just a heuristic guess
                error_lines.append("WARNING: The code is storing to the low address 0x{:x}".format(addr))
                error_lines.append("  The address contains code instead of data, which you should not be modifying.")
            elif call_states and addr <= 0x10000 and addr >= call_states[-1].sp:
                error_lines.append("WARNING: The code is storing to the stack address 0x{:x},".format(addr))
                error_lines.append("  which is above where the stack pointer was when the most recent function was called.")
                error_lines.append("  You may be modifying space allocated by a different function on the stack.")
        elif self.opcode == Riscv.opcode_JAL:
            newPC = self.getPC + self.jimm20
            if self.rd == 1:
                call_states.append(CallState(self, newPC))
        elif self.opcode == Riscv.opcode_JALR:
            newPC = self.getRegister(self.rs1) + self.oimm12
            if self.rs1 == 1 and self.oimm12 == 0:
                # return
                if call_states:
                    error_lines.extend(check_cc(call_states.pop(), CallState(self, newPC)))
                else:
                    error_lines.append("WARNING: The code is returning more times than functions have been called.")

        if total_error_lines[0] < 50:
            for line in error_lines:
                print("[PC = 0x{:04x}] {}".format(self.getPC, line))
            total_error_lines[0] += len(error_lines)
            if total_error_lines[0] >= 50:
                print("Too many errors, debugging output truncated")

        orig_execute(self)
    return execute
Riscv.RVMachine.execute = execute_wrapper(Riscv.RVMachine.execute)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Run the program on a simulated RISC-V machine.')
    parser.add_argument("program_name", nargs="?", type=str, default="", help='Program name. There should exist a corresponding .vmh file.')
    parser.add_argument("test_number", nargs="?", type=int, default=0, help='Test number, from 1 to 5.')
    parser.add_argument("--auto", action="store_true", help='AUTOGRADER use only')
    args = parser.parse_args()

    w.setAuto(args.auto)
    w.loadHex(args.program_name + ".vmh")
    w.loadArg(args.test_number)
    w.run(args.program_name)
    print(w.showStats())