#!/usr/bin/env python3

# step through an executable under cosimulation and report errors

#import objgraph

import sys
import os
import json
import codecs
import traceback
import multiprocessing
import pickle

libdir = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'lib')
sys.path.append(libdir)

import utils
import msp_fr5969_model as model
from mspdebug_driver import MSPdebug
from msp_emulator import Emulator
from msp_cosim import Cosim
from msp_isa import isa
import cosim as cosim_repl
import smt

def is_timer_read(fields, rdst):
    return fields['words'] == [0x4210 | rdst, 0x0350]

def is_reg_store(fields, rsrc):
    words = fields['words']
    if len(words) != 2 or words[0] != 0x40c2 | rsrc << 8:
        return -1
    return words[1]

def is_reg_sub(fields, rsrc, rdst):
    return fields['words'] == [0x8000 | rsrc << 8 | rdst]

run_max_steps = 20000
run_interval = 1
run_passes = 5

def check_elf(elfname, verbosity = 0):
    mulator = Emulator(verbosity=verbosity, tracing=True)
    mulator.prog(elfname)

    fram_end = mulator.md(model.upper_start - 256, 256)
    for byte in fram_end[:-2]:
        if byte != 255:
            print('Invalid prog write to reserved fram region:')
            print(utils.triple_summarize(fram_end, model.upper_start - 256))
            return False
    resetvec = fram_end[-2:]

    for i in range(run_passes):
        success, steps = mulator.run(run_max_steps)
        if not success:
            return False
        elif steps < run_max_steps:
            fram_end = mulator.md(model.upper_start - 256, 256)
            for byte in fram_end[:-2]:
                if byte != 255:
                    print('Invalid run write to reserved fram region:')
                    print(utils.triple_summarize(fram_end, model.upper_start - 256))
                    return False
            if fram_end[-2:] != resetvec:
                print('Broke the reset vector?')
                print('was', resetvec)
                print('now', fram_end[-2:])
                return False
            upper = mulator.md(model.upper_start, model.upper_size)
            for byte in upper:
                if byte != 255:
                    print('Modified upper memory?')
                    print(utils.triple_summarize(upper, model.upper_start))
                    return False
            touchdown = mulator.md(mulator.regs()[0], 16)
            expected_touchdown = [0xff, 0x3f, 0xff, 0x3f, 0xff, 0x3f, 0xff, 0x3f, 0xff, 0x3f, 0xff, 0x3f, 0xff, 0x3f, 0xf8, 0x3f]
            if touchdown != expected_touchdown:
                print('Missed touchdown pad:')
                print('  expecting', expected_touchdown)
                print('  got', touchdown)
                return False
            if verbosity >= 0:
                print('  checked {:s}, pass'.format(elfname))
            return True

    print('Did not complete?')
    return False

def trace_elf(elfname, jname, tty = None, logname = None, verbosity = 0):
    if logname is None:
        with MSPdebug(tty=tty, logf=sys.stdout, verbosity=verbosity) as driver:
            mulator = Emulator(verbosity=verbosity, tracing=True)
            mmap = [(model.ram_start, model.ram_size), (model.fram_start, model.fram_size)]
            cosim = Cosim([driver, mulator], [True, False], mmap)
            master_idx = 0

            cosim_repl.prog_and_sync(cosim, master_idx, elfname)
            cosim.run(max_steps=run_max_steps, interval=run_interval, passes=run_passes)

            diff = cosim.diff()
            trace = mulator.trace
            iotrace = mulator.iotrace2
    else:
        with open(logname, 'at') as f:
            with MSPdebug(tty=tty, logf=f, verbosity=max(verbosity,1)) as driver:
                mulator = Emulator(verbosity=verbosity, tracing=True)
                mmap = [(model.ram_start, model.ram_size), (model.fram_start, model.fram_size)]
                cosim = Cosim([driver, mulator], [True, False], mmap)
                master_idx = 0

                cosim_repl.prog_and_sync(cosim, master_idx, elfname)
                cosim.run(max_steps=run_max_steps, interval=run_interval, passes=run_passes)

                diff = cosim.diff()
                trace = mulator.trace
                iotrace = mulator.iotrace2

    with utils.Write7z(jname) as f:
        writer = codecs.getwriter('utf-8')
        json.dump({'diff':diff, 'trace':trace, 'iotrace':iotrace}, writer(f))

    if verbosity >= 0:
        print('  traced {:s} to {:s}'.format(elfname, jname))

def retrace_elf(elfname, jname, tinfo, interesting_blocks, verbosity = 0):
    if not os.path.isfile(jname):
        print('skipping {:s}, no trace {:s}'.format(elfname, jname))
        return True

    timulator = Emulator(verbosity=verbosity, tracing=True, tinfo=tinfo)
    mulator = Emulator(verbosity=verbosity, tracing=True)
    mmap = [(model.ram_start, model.ram_size), (model.fram_start, model.fram_size)]
    cosim = Cosim([timulator, mulator], [False, False], mmap)
    master_idx = 0

    cosim_repl.prog_and_sync(cosim, master_idx, elfname)
    cosim.run(max_steps=run_max_steps, interval=run_interval, passes=run_passes)

    tmp_jstr = json.dumps({'diff':cosim.diff(), 'trace':mulator.trace, 'iotrace':mulator.iotrace2})
    tmp_jobj = json.loads(tmp_jstr)

    diff = tmp_jobj['diff']
    trace = tmp_jobj['trace']
    iotrace = tmp_jobj['iotrace']
    old_diff, old_trace, old_iotrace = load_trace(jname)

    same = diff == old_diff
    if verbosity >= 0:
        print('  timed emulated {:s} against {:s}. Same? {:s}'
              .format(elfname, jname, repr(same)))
        # print('---ORIGINAL---')
        # utils.explain_diff(old_diff)
        # print('---EMULATOR---')
        # utils.explain_diff(diff)

    if not same:
        old_blocks = []
        old_mismatches = compute_mismatches(old_diff, verbosity=verbosity)
        old_err = mismatches_to_blocks(old_trace, old_mismatches, old_blocks)
        blocks = []
        mismatches = compute_mismatches(diff, verbosity=verbosity)
        err = mismatches_to_blocks(trace, mismatches, blocks)

        if old_err and err:
            print('    failures in both traces: {:s}'.format(elfname))
        elif old_err:
            print('    BAD: failures in hardware trace: {:s}'.format(elfname))
        elif err:
            print('    BAD: failures in emulator trace: {:s}'.format(elfname))
        else:
            print('    successful trace: {:s}'.format(elfname))

        old_blocks_index = {addr: (x, y) for (addr, x, y) in old_blocks}
        trace_errors = 0
        uncovered = 0
        for (addr, block, difference) in blocks:
            if addr in old_blocks_index:
                old_block, old_difference = old_blocks_index.pop(addr)
                if block != old_block:
                    print('      BAD: trace difference at {:05x}'.format(addr))
                    trace_errors += 1
                elif difference != old_difference:
                    interesting_blocks.append((addr, old_block, old_difference))
                    if verbosity >= 0:
                        print('timing difference for block at {:05x} of {:s}'.format(addr, elfname))
                        for fields in block:
                            ins = isa.decode(fields['words'][0])
                            fmt, name, smode, dmode = isa.instr_to_modes(ins)
                            if   fmt == 'fmt1':
                                rsrc = fields['rsrc']
                                rdst = fields['rdst']
                                if 'isrc' in fields:
                                    sval = ', {:#x}'.format(fields['isrc'])
                                else:
                                    sval = ''
                                print('{:s}\t{:s} (R{:d}{:s}), {:s} (R{:d})'
                                      .format(name, smode, rsrc, sval, dmode, rdst))
                            elif fmt == 'fmt2':
                                rsrc = fields['rsrc']
                                if 'isrc' in fields:
                                    sval = ', {:#x}'.format(fields['isrc'])
                                else:
                                    sval = ''
                                print('{:s}\t{:s} (R{:d}{:s})'
                                      .format(name, smode, rsrc, sval))
                            elif fmt == 'jump':
                                print('{:s}\t{:d}, taken={:s}'
                                      .format(name, fields['jump_offset'], str(fields['jump_taken'])))
                            else:
                                print('{:s}, {:s}, {:s}, {:s}'.format(fmt, name, smode, dmode))
                                utils.print_dict(fields)
                        print('hardware: {:s}, emulator: {:s}'
                              .format(repr(old_difference), repr(difference)))
                        print('')
            else:
                uncovered += 1

        if trace_errors > 0:
            print('    BAD: {:d} trace differences'.format(trace_errors))
        if uncovered > 0 or len(old_blocks_index) > 0:
            print('    BAD: {:d} blocks unique to hardware, {:d} to emulator'
                  .format(len(old_blocks_index), uncovered))

    return same

def load_trace(jname):
    with utils.Read7z(jname) as f:
        reader = codecs.getreader('utf-8')
        jobj = json.load(reader(f))
    return jobj['diff'], jobj['trace'], jobj['iotrace']

trace_suffix = '.trace.json.7z'

def compute_mismatches(diff, verbosity):
    mismatches = {}
    for addr in diff:
        if addr == 'regs':
            if verbosity >= 1:
                print('nonempty reg diff???')
                utils.explain_diff(diff)
        else:
            mems = diff[addr]
            assert len(mems) == 2
            assert len(mems[0]) == len(mems[1])
            for i in range(len(mems[0])):
                if mems[0][i] != mems[1][i]:
                    mismatches[int(addr)+i] = (mems[0][i], mems[1][i])
    return mismatches

# Consumes mismatches to add to blocks.
# Asserts that mismatches is empty when it finishes (no mismatches were not able
# to be explained by looking at the trace).
def mismatches_to_blocks(trace, mismatches, blocks):
    current = []
    in_region = False
    in_store = False
    missing = 0
    err = False
    for fields in trace:
        if is_timer_read(fields, 14):
            assert current == [] and in_region == False
            in_region = True
        elif is_timer_read(fields, 15):
            assert len(current) > 0 and in_region == True
            in_region = False
        elif is_reg_sub(fields, 14, 15):
            assert in_region == False
            in_store = True

        if in_store:
            assert in_region == False
            addr = is_reg_store(fields, 15)
            if addr > 0:
                if addr in mismatches:
                    blocks.append((addr, current, mismatches.pop(addr)))
                else:
                    # print('MISSING {:s} | {:s}'.format(repr(addr), repr(current)))
                    missing += 1
                    err = True
                current = []
                in_store = False

        elif in_region:
            current.append(fields)

    if err:
        print('Unexpected blocks! {:d} blocks have no corresponding mismatches, {:d} unexplained mismatches in diff.'
              .format(missing, len(mismatches)))

    elif len(mismatches) > 0:
        err = True
        print('State mismatch! {:d} unexplained mismatches in diff.'
              .format(len(mismatches)))
        print(mismatches)

    return err

def arff_header():
    s = "@relation 'cycle_count'\n"
    for i, ins in enumerate(isa.ids_ins):
        s += '@attribute {:s} numeric\n'.format('_'.join(isa.idx_to_modes(i)))
    return s + '@attribute cycles numeric\n@data'

def arff_entry(indices, cycles):
    bins = {}
    for k in indices:
        bins[k] = bins.setdefault(k, 0) + 1
    return ', '.join([str(bins[i]) if i in bins else '0' for i in range(len(isa.ids_ins))] + [str(cycles)])

def create_arff(blocks, arffname):
    with open(arffname, 'wt') as f:
        f.write(arff_header() + '\n')

        for addr, block, difference in blocks:
            assert(len(difference) == 2 and difference[1] == 0)
            cycles = difference[0]
            indices = []
            description = []
            for fields in block:
                ins = isa.decode(fields['words'][0])
                idx = isa.instr_to_idx(ins)
                modes = isa.idx_to_modes(idx)
                indices.append(idx)
                description.append('  {:s} {:s}, {:s}'.format(*modes[1:]))

            f.write(arff_entry(indices, cycles) + '\n')

def create_json(blocks, jname):
    with open(jname, 'wt') as f:
        json.dump({'blocks':blocks}, f)

def extract_json(jname):
    if jname.endswith('.json.7z'):
        with utils.Read7z(jname) as f:
            reader = codecs.getreader('utf-8')
            jobj = json.load(reader(f))
        return jobj['blocks']
    else:
        with open(jname, 'rt') as f:
            jobj = json.load(f)
        return jobj['blocks']

def walk_par(fn, targetdir, cargs, n_procs = 1, verbosity = 0):
    if os.path.isfile(targetdir):
        if verbosity >= 0:
            print('target is a single file, processing directly')
            return [fn((0, [('', targetdir)], cargs))]

    roots = set()
    worklists = [(i, [], cargs) for i in range(n_procs)]
    i = 0
    for root, dirs, files in os.walk(targetdir, followlinks=True):
        if root in roots:
            del dirs[:]
            continue
        else:
            roots.add(root)
        for fname in files:
            worklists[i%n_procs][1].append((root, fname))
            i += 1

    if verbosity >= 0:
        if n_procs == 1: