import sys
import re
import struct

op_names = ["MOVI",
            "MOVR",
            "GETM",
            "PUTM",
            "ADDI",
            "ADDR",
            "SUBI",
            "SUBR",
            "XORI",
            "XORR",
            "NOTI",
            "NOTR",
            "MULI",
            "MULR",
            "DIVI",
            "DIVR",
            "PUSH",
            "POOP",
            "CALL",
            "HALT",
            "NOPE"]

reg_names = ["R0", "R1", "R2", "R3", "S0", "S1", "S2", "S3", "IP", "BP", "SP"]
section_names = ["DATA:", "CODE:", "STACK:"]
section_flags = {s: i + 1 for i, s in enumerate(section_names)}
ops = {s: i for i, s in enumerate(op_names)}
regs = {s: i for i, s in enumerate(reg_names)}

assembled = bytearray()


def to_uint8(data):
    alphanum = re.compile("^[0-9]+$")
    if isinstance(data, int):
        return struct.pack("<B", data)
    elif data.startswith("0x"):
        return struct.pack("<B", int(data, 16))
    elif alphanum.match(data):  # only numbers
        return struct.pack("<B", int(data))
    else:
        return None


def to_uint16(data):
    alphanum = re.compile("^[0-9]+$")
    if isinstance(data, int):
        return struct.pack("<H", data)
    elif data.startswith("0x"):
        return struct.pack("<H", int(data, 16))
    elif alphanum.match(data):  # only numbers
        return struct.pack("<H", int(data))
    else:
        return None


def is_reg(data):
    if data not in reg_names:
        return False
    return True


def movi(op, dst, src):
    global assembled
    if (is_reg(dst)):
        op_val = to_uint8(ops[op])
        dst_val = to_uint8(regs[dst])
        src_val = to_uint16(src)
        if op_val and dst_val and src_val:
            assembled += op_val + dst_val + src_val
        else:
            sys.stderr.write(
                "ERROR WHILE ASSEMBLING UNKNOWN VALUES\n")
            return False
    else:
        sys.stderr.write(
            "ERROR WHILE ASSEMBLING UNKNOWN REGISTER: {}\n".format(dst))
        return False
    return True


def assemble_code(line):
    global assembled
    sys.stdout.write("CODE: ")
    instruction = [x for x in re.split('\W', line) if x]
    op_name = instruction[0]

    if op_name not in op_names:
        sys.stderr.write(
            "ERROR WHILE ASSEMBLING UNKNOWN OPERATION: {}\n".format(op_name))
        return False
    sys.stdout.write("{} {}\n".format(op_name, ", ".join(instruction[1::])))
    if op_name == "MOVI":
        movi(op_name, instruction[1], instruction[2])
    return True


def assemble_data(line):
    sys.stdout.write("DATA:\t")
    sys.stdout.write(line.strip(",") + "\n")


def main():
    global assembled
    if (len(sys.argv) < 2):
        print("Usage: {} <file_to_assemble>".format(sys.argv[0]))
        return

    with open(sys.argv[1], 'r') as f:
        gen = (line.strip("\n") for line in f if line != "\n")
        flag = None

        for line in gen:
            if line.startswith(tuple(section_names)):
                flag = section_flags[line]
                continue
            if flag == section_flags["DATA:"]:
                assemble_data(line)
            elif flag == section_flags["CODE:"]:
                assemble_code(line)
        if not flag:
            sys.stderr.write("Nothing was assembled! Did you use the section delimiters?\n")
    with open("./out.gipu", 'wb') as f:
        f.write(assembled)

if __name__ == '__main__':
    main()