#!/usr/bin/env python3
import re, csv, yaml, pathlib, collections

try:
    def dr(h,l): return list(range(h,l-1,-1))
    base = pathlib.Path('riscv-opcodes')
    opcodes = yaml.safe_load(open(base / 'instr_dict.yaml'))
    csrs = dict((int(a, 16), n) for fn in ['csrs.csv', 'csrs32.csv'] for a, n in csv.reader(open(base / fn), skipinitialspace=True))
    arg_bits = dict((a, dr(int(h),int(l))) for a, h, l in csv.reader(open(base / 'arg_lut.csv'), skipinitialspace=True))
    for s in open(base / 'constants.py').readlines():  # immediate scrambling from latex_mapping. Some better way?
        if mask := re.match(r"latex_mapping\[['\"](.*?)['\"]\] = ['\"][^\[]*\[([^\]]*)\]['\"]", s):
            fbits = sum([(dr(*(int(i) for i in part.split(':'))) if ':' in part else [int(part)]) for part in mask[2].split('$\\\\vert$')], [])
            locs = [-1] * (max(fbits)+1)
            for i, b in enumerate(fbits): locs[-b-1] = arg_bits[mask[1]][i]
            arg_bits[mask[1]] = locs #[31] * (32-len(locs)) + locs if locs[0] == 31 else locs  # sign extension to 32 bits
except Exception as e: raise Exception("Unable to load RISC-V specs. Do:\n"
                                       "git clone https://github.com/riscv/riscv-opcodes.git\n"
                                       "make -C riscv-opcodes")

for name, op in opcodes.items():
    op['name'], op['mask'], op['match'] = name, int(op['mask'], 16), int(op['match'], 16)
    del op['encoding']
    op['arg_bits'] = {}
    for vf in op['variable_fields']:
        if vf not in arg_bits: continue
        bits = [-1] * (32-len(arg_bits[vf])) + arg_bits[vf]  # padding for proper alignment when combining hi and lo fields
        vf2 = vf.replace('hi','').replace('lo','').replace('c_','')
        op['arg_bits'][vf2] = [max(a,b) for a, b in zip(op['arg_bits'][vf2], bits)] if vf2 in op['arg_bits'] else bits
    for vf in op['arg_bits']: op['arg_bits'][vf] = op['arg_bits'][vf][next(n for n, e in enumerate(op['arg_bits'][vf]) if e >= 0):]  # remove padding

common_ops = ('addi,sw,lw,jal,bne,beq,add,jalr,lbu,slli,lui,andi,or,bltu,srli,and,sub,blt,bgeu,xor,sb,auipc,sltiu,bge,lb,mul,sltu,lhu,sll,srl,sh,amoadd_w,xori,ori,csrrci,csrrs,srai,fence,lr_w,sc_w,mulhu,amoor_w,lh,amoand_w,csrrsi,divu,div,remu,sra,slt,csrrw,amoswap_w,csrrc,mulh,fence_i,rem,mret,csrrwi').split(',')

mask_match = []
mask_match_aliases = collections.defaultdict(set)
for mask in dict((opcodes[op]['mask'],1) for op in common_ops + list(opcodes)):
    matches = {}
    for op in opcodes.values():
        if op['mask'] != mask or op['name'].endswith('_rv32'): continue
        if op['match'] in matches: mask_match_aliases[matches[op['match']][1:-1]].add(op['name'])
        else: matches[op['match']] = f'${op["name"]}$'
    mask_match.append((mask, matches))

print('writing tinyrv/opcodes.py')
with open('tinyrv/opcodes.py', 'w') as f:
    f.write(f'# auto-generated by tinyrv_opcodes_gen.py\n')
    f.write(f'opcodes={str(opcodes)}\n')
    f.write(f'arg_bits={str(arg_bits)}\n')
    f.write(f'csrs={str(csrs)}\n')
    dstr = str(mask_match).replace("'$", "opcodes['").replace("$'", "']")
    f.write(f'mask_match={dstr}\n')
    f.write(f'mask_match_aliases={dict(mask_match_aliases)}\n')