#!/usr/bin/env python3
# a simple assembler/disassembler for the n808 VM
# supports both text and binary program formats
# usage: n8asm [at|ab|dt|db|t2b|b2t|t74] file outfile
# modes:
# at: assemble to plaintext (N8) machine code
# ab: assemble to binary (N8B) machine code
# dt: disassemble plaintext (N8) machine code
# db: disassemble binary (N8B) machine code
# t2b: convert from plaintext N8 to binary N8B
# b2t: convert from binary N8B to plaintext N8
# t74: convert from plaintext N8 to the TI-74 representation
# Created by Luxferre in 2025, released into public domain

import sys, re, struct

# n808 mnemonics for assembly and disasssembly
mnemos = ['nop', 'jmp', 'iat', 'ino', 'cpy', 'set', 'mat', 'rnd']

# n808 jump/function shortcuts (arranged from longest to shortest)
jmpfunc_shorts = {
    'ret': 'jmp 13 0 125', # return from a procedure 
    'nnn': 'nop 0 0 0',    # alias for nop
    'inc': 'mat 0 127',    # increment
    'dec': 'mat 0 126',    # decrement
    'inv': 'mat 3 127',    # inverse/reciprocal
    'iuc': 'jmp 13 0',     # indirect unconditional jump
    'jpr': 'jmp 14 0',     # jump to a procedure
    'neg': 'mat 1 0',      # negation
    'juc': 'jmp 6 0',      # direct unconditional jump
    '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
    'cos': 'mat 10',       # cosine
    'atn': 'mat 11',       # arctangent
    '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
    '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
    'out': 'ino 0',        # normal numeric output
    'inp': 'ino 1',        # normal numeric input
    'ouc': 'ino 2',        # character output
    'ipc': 'ino 3',        # character input
    'dca': 'cpy 0',        # direct constant assignment
    'dva': 'cpy 1',        # direct value assignment
    'ica': 'cpy 2',        # indirect constant assignment
    'iva': 'cpy 3',        # indirect value assignment
    'ivc': 'cpy 4',        # indirect value copy
    'add': 'mat 0',        # addition
    'sub': 'mat 1',        # subtraction
    'mul': 'mat 2',        # multiplication
    'div': 'mat 3',        # division
    'mdf': 'mat 4',        # modulo/floor
    'abs': 'mat 5',        # absolute value
    'sqr': 'mat 6',        # square root
    'exp': 'mat 7',        # natural exponent
    'log': 'mat 8',        # natural logarithm
    'sin': 'mat 9'         # sine
}

# converts the N8A source code to the plaintext machine code representation
def assemble(source):
    labels = {}
    out = ''
    lno = 0
    vreg = re.compile(r'\s+')
    for shrt, meaning in jmpfunc_shorts.items(): # replace the 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
            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
            elif fields[0].startswith('#'): # this is an alias
                labels['@' + fields[1].strip()] = int(fields[0][1:])
                continue # do not update the instruction number
            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(cmd), fields[1], fields[2], fields[3]]) + '\n'
            except IndexError:
                print('Assembly error - not enough operands at line', lno)
                sys.exit(1)
            lno += 1 # update the instruction number
    for lbl, lineno in labels.items():
        out = out.replace(lbl, str(lineno))
    final = '' # prepare the final text
    for line in out.split('\n'):
        if len(line) > 0:
            fields = vreg.split(line)[:4] # split it into fields
            instr = int(fields[0]) * 2097152 + int(fields[1]) * 16384
            instr += int(fields[2]) * 128 + int(fields[3])
            final += str(instr) + '\n'
    return final

# converts the plaintext machine code representation to N8A source code
def disassemble(machcode):
    vreg = re.compile(r'\s+')
    iindex = 0
    instrs = []
    for instrnum in vreg.split(machcode):
        if len(instrnum) > 0:
            instrnum = int(instrnum)
            opcode = (instrnum >> 21) & 7
            p1 = (instrnum >> 14) & 127
            p2 = (instrnum >> 7) & 127
            p3 = instrnum & 127
            instrs.append([iindex, opcode, p1, p2, p3])
            iindex += 1
    # now, instrs contains 5-number groups
    assembly = [None for i in range(0, 128)]
    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
    for shrt, meaning in jmpfunc_shorts.items(): # replace the shortcuts
        out = out.replace(meaning, shrt)
    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+')
    out = b''
    for instr in vreg.split(txtrep):
        if len(instr) > 0:
            instr = int(instr)
            out += struct.pack('BBB', (instr >> 16) & 255, (instr >> 8) & 255, instr & 255)
    return out

# converts the binary machine code representation to plaintext representation
def totext(binrep):
    out = ''
    while len(binrep) > 0:
        chunk = binrep[0:3]
        binrep = binrep[3:]
        b1, b2, b3 = struct.unpack('BBB', chunk)
        instr = (b1 << 16) | (b2 << 8) | b3
        out += str(instr) + '\n'
    return out

# converts the N8 plaintext machine code to the TI-74 DATA statements (N74)
def ti74data(txtrep:str, baseaddr:int=1000):
    instrs = [] # instruction list to store here
    vreg = re.compile(r'\s+')
    out = ''
    iindex = 0
    for instr in vreg.split(txtrep):
        if len(instr) > 0:
            instrs.append(str(int(instr)))
    instrs.append(str(-1)) # add the terminating instruction
    while len(instrs) > 0:
        dchunk = instrs[:7]
        instrs = instrs[7:]
        out += f'{baseaddr + iindex} DATA {','.join(dchunk)}\n'
        iindex += 1
    return out

def main():
    if len(sys.argv) < 4:
        print("Usage: n8asm [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()