import os
import sys
import traceback

sys.path.insert(1, os.path.join(sys.path[0], '..'))

import infernal

################################################################################

tests = []


def add_test(name, result, register, code):
    tests.append((name, result, register, code))


def printf(str, *args):
    print(str.format(*args))


def execute_tests():
    print('Executing tests!')
    total_tests = 0
    failed_tests = 0
    error_tests = 0
    for name, result, register, code in tests:
        total_tests += 1
        line_nr = None
        try:
            emu = infernal.Emulator(code)
            emu.setStack('junk...', 'calling eip')
            emu.setRegs(rip=0, rbp='old bp')
        except infernal.CodeParseException as e:
            error_tests += 1
            printf(
                'Encountered error when parsing {}, at line {}: {}',
                name,
                e.line_nr,
                e.str,
            )
        try:
            for line_nr in emu:
                pass
            if isinstance(result, BaseException):
                printf('Error should have happened in {}, but did not', name)
                failed_tests += 1
            output = emu.getVal(register)
            if output != result:
                failed_tests += 1
                printf(
                    'Failed in {}. {} was {}, should be {}',
                    name,
                    register,
                    output,
                    result,
                )
        except BaseException as e:
            print(e)
            if not isinstance(result, BaseException) or not isinstance(e, result):
                error_tests += 1
                printf(
                    'Encountered error in {}, at operation {}: {}',
                    name,
                    emu.getVal('%rip'),
                    e,
                )
                traceback.print_exc()
    printf('Tests done! {}/{}.', total_tests - failed_tests - error_tests, total_tests)
    printf('{} failed, and {} encountered a python error', failed_tests, error_tests)


################################################################################
# Arithmetic Operations

add_test(
    'constant $255',
    255,
    '%rsi',
    """
movq $255, %rsi
""",
)

add_test(
    'static addition 10+$20',
    30,
    '%rsi',
    """
movq $10, %rsi
addq $20, %rsi
""",
)

add_test(
    'register addition 10+20',
    30,
    '%rsi',
    """
movq $10, %rsi
movq $20, %rax
addq %rax, %rsi
""",
)

add_test(
    'register subtraction 10-20',
    -10,
    '%rsi',
    """
movq $10, %rsi
movq $20, %rax
subq %rax, %rsi
""",
)

################################################################################
# Branching

branch_tests = [
    ('jg', 10, '>', 5, 0),
    ('jg', 5, '>', 5, 1),
    ('jg', 5, '>', 10, 1),
    ('jl', 10, '<', 5, 1),
    ('jl', 5, '<', 5, 1),
    ('jl', 5, '<', 10, 0),
    ('je', 10, '==', 5, 1),
    ('je', 5, '==', 5, 0),
    ('je', 5, '==', 10, 1),
    ('jge', 10, '>=', 5, 0),
    ('jge', 5, '>=', 5, 0),
    ('jge', 5, '>=', 10, 1),
    ('jle', 10, '<=', 5, 1),
    ('jle', 5, '<=', 5, 0),
    ('jle', 5, '<=', 10, 0),
    ('jne', 10, '!=', 5, 0),
    ('jne', 5, '!=', 5, 1),
    ('jne', 5, '!=', 10, 0),
]

for jump_instruct, a, comp, b, result in branch_tests:
    add_test(
        f'branch {a} {comp} {b}={not result}',
        result,
        '%rsi',
        f"""
    start:	cmpq ${b}, ${a}
            {jump_instruct} true
            movq $1, %rsi
            jmp return
    true:	movq $0, %rsi
    return:	ret
    """,
    )

################################################################################
# Junk Comparisons

add_test(
    'invalid comparison 1',
    infernal.JunkComparisonException,
    'None',
    """
start:	movq $100, %rsp   # Set stack pointer to a random position.
        popq %rsi		  # Move a Junk value into %rsi
        cmpq $10, %rsi    # Do a Junk comparison, which triggers the `i` comp
                          # virtual register.
        jg start		  # Attempt to do a jump to the start, if (Junk-10)>0,
                          # which makes no sense, and thus throws an error.
""",
)

add_test(
    'invalid addition 1',
    infernal.JunkComparisonException,
    'None',
    """
start:	movq $100, %rsp   # Set stack pointer to a random position.
        popq %rsi		  # Move a Junk value into %rsi
        addq %rsi, %rsp   # Adds Junk to 101, which produces Junk.
        cmpq $10, %rsp    # Do a Junk comparison, which triggers the `i` comp
                          # virtual register.
        jg start		  # Attempt to do a jump to the start, if (Junk-10)>0,
                          # which makes no sense, and thus throws an error.
""",
)

add_test(
    'invalid addition 2',
    infernal.JunkComparisonException,
    'None',
    """
start:	movq $100, %rsp   # Set stack pointer to a random position.
        popq %rsi		  # Move a Junk value into %rsi
        addq %rsp, %rsi   # Adds 101 to Junk, which produces Junk.
        cmpq $10, %rsi    # Do a Junk comparison, which triggers the `i` comp
                          # virtual register.
        jg start		  # Attempt to do a jump to the start, if (Junk-10)>0,
                          # which makes no sense, and thus throws an error.
""",
)

add_test(
    'invalid subtraction 1',
    infernal.JunkComparisonException,
    'None',
    """
start:	movq $100, %rsp   # Set stack pointer to a random position.
        popq %rsi		  # Move a Junk value into %rsi
        subq %rsi, %rsp   # Adds Junk to 101, which produces Junk.
        cmpq $10, %rsp    # Do a Junk comparison, which triggers the `i` comp
                          # virtual register.
        jg start		  # Attempt to do a jump to the start, if (Junk-10)>0,
                          # which makes no sense, and thus throws an error.
""",
)

add_test(
    'invalid subtraction 2',
    infernal.JunkComparisonException,
    'None',
    """
start:	movq $100, %rsp   # Set stack pointer to a random position.
        popq %rsi		  # Move a Junk value into %rsi
        subq %rsp, %rsi   # Adds 101 to Junk, which produces Junk.
        cmpq $10, %rsi    # Do a Junk comparison, which triggers the `i` comp
                          # virtual register.
        jg start		  # Attempt to do a jump to the start, if (Junk-10)>0,
                          # which makes no sense, and thus throws an error.
""",
)

################################################################################

if __name__ == '__main__':
    execute_tests()