#!/usr/bin/env python3
# a simple assembler/disassembler for the mu808 VM
# supports both text and binary program formats
# usage: mu808asm [at|ab|dt|db|t2b|b2t|t74] file outfile
# modes:
# at: assemble to plaintext (MU8) machine code
# ab: assemble to binary (MU8B) machine code
# dt: disassemble plaintext (MU8) machine code
# db: disassemble binary (MU8B) machine code
# t2b: convert from plaintext MU8 to binary MU8B
# b2t: convert from binary MU8B to plaintext MU8
# t74: convert from plaintext MU8 to the TI-74 representation

import sys, re, struct

# mu808 mnemonics for assembly and disasssembly
mnemos = ['nop', 'jmp', 'iat', 'out', 'inp', 'set', 'cpy', 'fma',
         'sub', 'div', 'mdf', 'abs', 'sqr', 'nel', 'tri', 'rnd']

# mu808 jump shortcuts
jmp_shorts = {
    'jeq': 'jmp 0',    # direct jump if equals to zero
    'jgt': 'jmp 1',    # direct jump if greater than zero
    'jlt': 'jmp 2',    # direct jump if less than zero
    'jge': 'jmp 3',    # direct jump if greater than or equals to zero
    'jle': 'jmp 4',    # direct jump if less than or equals to zero
    'jne': 'jmp 5',    # direct jump if not equals to zero
    'juc': 'jmp 6 0',  # direct unconditional jump
    'ieq': 'jmp 7',    # indirect jump if equals to zero
    'igt': 'jmp 8',    # indirect jump if greater than zero
    'ilt': 'jmp 9',    # indirect jump if less than zero
    'ige': 'jmp 10',   # indirect jump if greater than or equals to zero
    'ile': 'jmp 11',   # indirect jump if less than or equals to zero
    'ine': 'jmp 12',   # indirect jump if not equals to zero
    'iuc': 'jmp 13 0'  # indirect unconditional jump
}

# converts the MU8A source code to the plaintext machine code representation
def assemble(source):
    labels = {}
    out = ''
    lno = 0
    vreg = re.compile(r'\s+')
    for shrt, meaning in jmp_shorts.items(): # replace the jump shortcuts
        source = source.replace(shrt, meaning)
    for line in source.split('\n'): # parse the main code
        fields = []
        line = line.split(';')[0].strip()
        if len(line) > 0: # actual line to be counted on
            lno += 1 # populate the line number
            fields = vreg.split(line)[:5] # split it into fields
            if fields[0].startswith(':'): # this is a label
                labels[fields[0]] = lno # remember the label
                fields = fields[1:] # remove the label from fields
            try:
                cmd = mnemos.index(fields[0].lower())
            except ValueError:
                try:
                    cmd = int(cmd)
                except:
                    print('Assembly error - unknown mnemonic at line', lno)
                    sys.exit(1)
            try:
                out += ' '.join([str(lno), str(cmd), fields[1], fields[2], fields[3]]) + '\n'
            except IndexError:
                print('Assembly error - not enough operands at line', lno)
                sys.exit(1)

    for lbl, lineno in labels.items():
        out = out.replace(lbl, str(lineno))
    return out

# converts the plaintext machine code representation to MU8A source code
def disassemble(machcode):
    vreg = re.compile(r'\s+')
    iindex = 0
    instr = []
    instrs = []
    for instrpart in vreg.split(machcode):
        if iindex == 5:
            instrs.append(instr)
            instr = []
            iindex = 0
        if len(instrpart) > 0:
            instr.append(int(instrpart))
        iindex += 1
    # now, instrs contains 5-number groups
    assembly = [None for i in range(0, 16384)]
    labels = {}
    for instr in instrs: # iterate over each instruction
        lno, cmd, x, y, z = instr
        try:
            mnemo = mnemos[cmd]
        except IndexError:
            print('Disassembly error - unknown opcode at line', lno)
            sys.exit(1)
        if cmd == 1: # save the label for jump instructions
            labels[z] = ':lbl_' + str(z)
            assembly[lno] = [mnemo, str(x), str(y), labels[z]]
        else:
            assembly[lno] = [mnemo, str(x), str(y), str(z)]
    out = ''
    lno = 0
    for asline in assembly:
        if lno in labels:
            out += labels[lno] + ' '
        if asline is not None:
            out += ' '.join(asline) + '\n'
        lno += 1
    return out

# converts the plaintext machine code representation to binary representation
def tobinary(txtrep):
    instrs = [] # instruction list to store here
    vreg = re.compile(r'\s+')
    iindex = 0
    instr = []
    for instrpart in vreg.split(txtrep):
        if iindex == 5:
            instrs.append(instr)
            instr = []
            iindex = 0
        if len(instrpart) > 0:
            instr.append(int(instrpart))
        iindex += 1
    # now, instrs contains 5-number groups
    out = b''
    # the format is: cmd (8 bits), lno (14 bits), arg1 to arg3 (14 bits each)
    for instr in instrs: # iterate over each instruction
        b1 = instr[1] & 15
        b2 = (instr[0] >> 6) & 255
        b3 = ((instr[0] << 2) | (instr[2] >> 12)) & 255
        b4 = (instr[2] >> 4) & 255
        b5 = ((instr[2] << 4) | (instr[3] >> 10)) & 255
        b6 = (instr[3] >> 2) & 255
        b7 = ((instr[3] << 6) | (instr[4] >> 8)) & 255
        b8 = instr[4] & 255
        out += struct.pack('BBBBBBBB', b1, b2, b3, b4, b5, b6, b7, b8)
    return out

# converts the binary machine code representation to plaintext representation
def totext(binrep):
    out = ''
    while len(binrep) > 0:
        chunk = binrep[0:8]
        binrep = binrep[8:]
        b1, b2, b3, b4, b5, b6, b7, b8 = struct.unpack('BBBBBBBB', chunk)
        cmd = b1 & 15
        lno = ((b2 << 6) | (b3 >> 2)) & 16383
        x = ((b3 << 12) | (b4 << 4) | (b5 >> 4)) & 16383
        y = ((b5 << 10) | (b6 << 2) | (b7 >> 6)) & 16383
        z = ((b7 << 8) | b8) & 16383
        out += ' '.join([str(lno), str(cmd), str(x), str(y), str(z)]) + '\n'
    return out

# converts the MU8 plaintext machine code to the TI-74 DATA statements (MU74)
def ti74data(txtrep:str, baseaddr:int=5000):
    instrs = [] # instruction list to store here
    vreg = re.compile(r'\s+')
    iindex = 0
    instr = []
    for instrpart in vreg.split(txtrep):
        if iindex == 5:
            instrs.append(instr)
            instr = []
            iindex = 0
        if len(instrpart) > 0:
            instr.append(int(instrpart))
        iindex += 1
    # now, instrs contains 5-number groups
    out = f'{baseaddr} DATA 0,0\n'
    maxlno = 0
    for instr in instrs:
        lno, cmd, x, y, z = instr
        if lno > maxlno:
            maxlno = lno
        a = cmd*10000 + x
        b = y*10000 + z
        out += f'{baseaddr + lno} DATA {a},{b}\n'
    out += f'{baseaddr + maxlno + 1} DATA -1,0\n'
    return out

def main():
    if len(sys.argv) < 4:
        print("Usage: mu808asm [at|ab|dt|db|t2b|b2t|t74] file outfile")
        return
    mode = sys.argv[1]
    srcfile = sys.argv[2]
    targetfile = sys.argv[3]
    srctext = ''
    inpmode = 'r'
    if mode == 'b2t' or mode == 'db':
        inpmode = 'rb'
    try:
        fd = open(srcfile, inpmode)
        srctext = fd.read()
        close(fd)
    except:
        pass
    if len(srctext) > 0:
        outmode = 'w'
        output = ''
        if mode == 'b2t': # binary-to-text
            output = totext(srctext)
        elif mode == 't2b': # text-to-binary
            output = tobinary(srctext)
            outmode = 'wb'
        elif mode == 'dt': # disassemble text
            output = disassemble(srctext)
        elif mode == 'db': # disassemble binary
            srctext = totext(srctext)
            output = disassemble(srctext)
        elif mode == 't74': # convert to TI-74 DATA statements
            output = ti74data(srctext)
        else: # assemble
            output = assemble(srctext) # default mode
            if mode == 'ab':
                output = tobinary(output)
                outmode = 'wb'
        # write the output file
        try:
            fd = open(targetfile, outmode)
            fd.write(output)
            close(fd)
        except:
            pass
    else:
        print('Nothing to process!')

if __name__ == '__main__':
    main()