from riscemu.colors import FMT_ERROR, FMT_NONE, FMT_BOLD, FMT_GREEN
from riscemu.instructions import InstructionSet
from riscemu.types import Instruction, CPU
from riscemu.decoder import RISCV_REGS

FMT_SUCCESS = FMT_GREEN + FMT_BOLD


def assert_equals(ins: Instruction, cpu: CPU):
    a, b = (get_arg_from_ins(ins, i, cpu) for i in (0, 2))
    return a == b


def assert_equals_mem(ins: Instruction, cpu: CPU):
    a, b = (get_arg_from_ins(ins, i, cpu) for i in (0, 2))
    a = cpu.mmu.read_int(a)
    return a == b


def assert_in(ins: Instruction, cpu: CPU):
    a = get_arg_from_ins(ins, 0, cpu)
    others = [get_arg_from_ins(ins, i, cpu) for i in range(2, len(ins.args))]
    return a in others


def _not(func):
    def test(ins: Instruction, cpu: CPU):
        return not func(ins, cpu)

    return test


def get_arg_from_ins(ins: Instruction, num: int, cpu: CPU):
    a = ins.args[num]
    if a in RISCV_REGS:
        return cpu.regs.get(a)
    return ins.get_imm(num)


assert_ops = {
    '==': assert_equals,
    '!=': _not(assert_equals),
    'in': assert_in,
    'not_in': _not(assert_in),
}


class Z_test(InstructionSet):
    def __init__(self, cpu: 'CPU'):
        print('[Test] loading testing ISA, this is only meant for running testcases and is not part of the RISC-V ISA!')
        self.failed = False
        super().__init__(cpu)

    def instruction_assert(self, ins: Instruction):
        if len(ins.args) < 3:
            print(FMT_ERROR + '[Test] Unknown assert statement: {}'.format(ins) + FMT_NONE)
            return
        op = ins.args[1]
        if op not in assert_ops:
            print(FMT_ERROR + '[Test] Unknown operation statement: {} in {}'.format(op, ins) + FMT_NONE)
            return

        if assert_ops[op](ins, self.cpu):
            print(FMT_SUCCESS + '[TestCase] 🟢 passed assertion {}'.format(ins))
        else:
            print(FMT_ERROR + '[TestCase] 🔴 failed assertion {}'.format(ins))
            self.cpu.halted = True
            self.failed = True

    def instruction_fail(self, ins: Instruction):
            print(FMT_ERROR + '[TestCase] 🔴 reached fail instruction! {}'.format(ins))
            self.cpu.halted = True
            self.failed = True