'''
The Overture CPU assembler.

Author: CoccaGuo
Date: 2023/07/15

Registers:
t0~t3 temp regs,
s0~s1 saved regs,
sp stack ptr,
io gpio.

Instructions:
 - OR
 - NAND
 - NOR
 - AND
 - ADD
 - SUB
 - EQ0
 - LT0
 - LTEQ0
 - JMP
 - NEQ0
 - GTEQ0
 - GT0
 - SW
 - LW

 - LABEL name
 - name

 - MOV R# R#
 - IMME number
 - #DEFINE name const
 - raw_cmd


 Use `;` for comments.
'''
REGS = {
    'T0': 0,
    'T1': 1,
    'T2': 2,
    'T3': 3,
    'S0': 4,
    'S1': 5,
#    'SP': 6,
    'IO': 6,
}

CMDS = {
    # CALC CMDS
    'OR':   0b01_000_000,
    'NAND': 0b01_000_001,
    'NOR':  0b01_000_010,
    'AND':  0b01_000_011,
    'ADD':  0b01_000_100,
    'SUB':  0b01_000_101,
    # COND CMDS
    'EQ0':  0b11_000_001,
    'LT0':  0b11_000_010,
    'LTEQ0':0b11_000_011,
    'JMP':  0b11_000_100,
    'NEQ0': 0b11_000_101,
    'GTEQ0':0b11_000_110,
    'GT0':  0b11_000_111,
    # MEM CMDS; NOT USED HERE
 #  'SW':   0b01_000_110,
 #  'LW':   0b01_000_111,
}


def digit(op: str):
    try:
        i = int(op, base=0)
    except ValueError:
        return None
    return i


def register(reg: str):
    if reg in REGS.keys():
        return REGS.get(reg)

    if reg.startswith('R'):
        reg = reg.strip('R')
    if reg.startswith('REG'):
        reg = reg.strip('REG')
    if reg.startswith('$'):
        reg = reg.strip('$')

    try:
        i = int(reg.strip('R'))
        if i>=0 and i<8:
            return i
        else:
            return None
    except ValueError:
        return None


def _macro(instr: str):
    '''
    convert defs into its defination.
    '''

    defs = dict()

    ops = instr.upper().split('\n')

    for ind, op in enumerate(ops):
        op = op.strip()
        if not op: continue
        if op.startswith(';'): continue

        if op.startswith('#'):
            op = op.strip('#').strip()
        if  op.startswith('DEFINE'):
            _defs = op.split()
            if len(_defs) < 3:
                raise SyntaxError(f'Line {ind+1}: `#DEFINE` should follow with an identifier and a const.')
            _name = _defs[1].strip()
            _const = ' '.join(_defs[2:]).strip()
            if not _name or not _const:
                raise SyntaxError(f'Line {ind+1}: `#DEFINE` should follow with an identifier and a const.')
            if _name in defs.keys():
                raise SyntaxError(f'Line {ind+1}: Multiple definition of #DEFINE `{_name}`.')
            defs[_name] = _const

    replaced = list()
    lines = instr.upper().split('\n')
    for l in lines:
        tokens = [defs.get(t, t) for t in l.split()]
        replaced.append(' '.join(tokens).strip())
    return '\n'.join(replaced).strip()


def _compile(instr: str):
    '''
    Compile asm into binarys.
    '''
    bin = list() # binary code
    labels = dict()

    pc = 0 # prog counter
    errs = list()

    ops = instr.upper().split('\n')

    for ind, op in enumerate(ops):
        op = op.strip()
        if not op: continue
        if op.startswith(';'): continue
        if op.startswith('#'): continue

        # parsing command
        if op in CMDS.keys():
             bin.append(CMDS.get(op))
             pc += 1
             continue

        # parsing label
        if op.startswith('LABEL'):
            if len(op.split('LABEL')) < 2:
                errs.append(SyntaxError(f'Line {ind+1}: `LABEL` should follow an identifier.'))
                continue
            _label = op.split('LABEL')[1].strip()
            if not _label:
                errs.append(SyntaxError(f'Line {ind+1}: `LABEL` should follow with an identifier.'))
            elif labels.get(_label):
                errs.append(SyntaxError(f'Line {ind+1}: Multiple definition of LABEL `{_label}`.'))
            else:
                labels[_label] = pc
                continue

        if op in labels.keys():
            _pc_ind = labels.get(op)
            if _pc_ind > 63:
                errs.append(IndexError(f'Line {ind+1}: LABEL index out of range (0~63).'))
            elif _pc_ind == 0x00:
                errs.append(IndexError(f'Line {ind+1}: LABEL index should not be zero.'))
            bin.append(labels.get(op))
            pc += 1
            continue


        # parse mov
        if op.startswith('MOV'):
            _regs = op.split()
            if len(_regs) != 3:
                errs.append(SyntaxError(f'Line {ind+1}: `MOV` should follow with two regs.'))
                continue
            _reg_src = register(_regs[1].strip())
            _reg_dst = register(_regs[2].strip())
            if _reg_src is not None and _reg_dst is not None:
                _cmd = 0x80 | _reg_src << 3 | _reg_dst
                bin.append(_cmd)
                pc += 1
                continue
            else:
                errs.append(SyntaxError(f'Line {ind+1}: `MOV` should follow with two regs.'))
                continue

        # parse imme
        if op.startswith('IMME'):
            _imme = op.replace('IMME', '').strip()
            _imme = digit(_imme)
            if _imme is None:
                errs.append(SyntaxError(f'Line {ind+1}: `IMME` should follow with an integer.'))
                continue
            if _imme >= 0 or _imme < 64:
                bin.append(_imme)
                pc += 1
                continue
            else:
                errs.append(SyntaxError(f'Line {ind+1}: Immediate number out of range (0~63).'))
                continue


        # parse raw
        _raw = digit(op)
        if _raw is None:
            errs.append(SyntaxError(f'Line {ind+1}: Unidentified instruction `{op}`.'))
        elif _raw > 255 or _raw < 0:
            errs.append(IndexError(f'Line {ind+1}: Command out of range 0~255.'))
        else:
            bin.append(_raw)
            pc += 1
            continue

    if not errs:
        return 0, bin
    else:
        return -1, errs


if __name__ == '__main__':
    code = r'''
    ; Example code for this tiny assembler,
    ; The Maze

    ; forward a step
    ; turn left
    ; if left is wall:
    ; turn right till not wall
    ; goto 1

    # define left 0
    # define forward 1
    # define right 2
    # define pass 3
    # define interact 4
    # define wall 1
    # define hlt 0b10_100_100

    hlt
    label start
    imme forward
    mov t0 io
    imme left
    mov t0 io
    label read_input
    mov io t3
    start
    eq0
    right
    mov t0 io
    interact
    mov t0 io
    read_input
    jmp
    '''
    asm = _macro(code)
    code, arr = _compile(asm)
    if code == 0:
        for b in arr:
            print(b, end=' ')
    else:
        for err in arr:
            print(err)
    print()