import os
import sys
import struct
import sqlite3
import argparse

version_no = 3

file_list = []
file_index_current = 0
object_index_current = 0

database_name = "klee_profiling.db"

""" Check if a database exists, otherwise create one """

if os.path.isfile(database_name):
    conn = sqlite3.connect(database_name)
    c = conn.cursor()
    print("Opened already created database")
else:
    print("No file named %s exists in current folder!")
    exit()


class KTestError(Exception):
    pass


class KTest:

    @staticmethod
    def fromfile(path):
        if not os.path.exists(path):
            print("ERROR: file %s not found" % (path))
            sys.exit(1)

        f = open(path, 'rb')
        hdr = f.read(5)
        if len(hdr) != 5 or (hdr != b'KTEST' and hdr != b"BOUT\n"):
            raise KTestError('unrecognized file')
        version, = struct.unpack('>i', f.read(4))
        if version > version_no:
            raise KTestError('unrecognized version')
        numArgs, = struct.unpack('>i', f.read(4))
        args = []
        for i in range(numArgs):
            size, = struct.unpack('>i', f.read(4))
            args.append(str(f.read(size).decode(encoding='ascii')))

        if version >= 2:
            symArgvs, = struct.unpack('>i', f.read(4))
            symArgvLen, = struct.unpack('>i', f.read(4))
        else:
            symArgvs = 0
            symArgvLen = 0

        numObjects, = struct.unpack('>i', f.read(4))
        objects = []
        for i in range(numObjects):
            size, = struct.unpack('>i', f.read(4))
            name = f.read(size)
            size, = struct.unpack('>i', f.read(4))
            bytes = f.read(size)
            objects.append((name, bytes))

        # Create an instance
        b = KTest(version, args, symArgvs, symArgvLen, objects)
        # Augment with extra filename field
        b.filename = path
        return b

    def __init__(self, version, args, symArgvs, symArgvLen, objects):
        self.version = version
        self.symArgvs = symArgvs
        self.symArgvLen = symArgvLen
        self.args = args
        self.objects = objects

        # add a field that represents the name of the program used to
        # generate this .ktest file:
        program_full_path = self.args[0]
        program_name = os.path.basename(program_full_path)
        # sometimes program names end in .bc, so strip them
        if program_name.endswith('.bc'):
            program_name = program_name[:-3]
        self.programName = program_name


def trimZeros(str):
    for i in range(len(str))[::-1]:
        if str[i] != '\x00':
            return str[:i + 1]

    return ''


def ktest_setdata(filename):
    """
    Substitute every variable found in ktest-file
    """
    global file_list
    b = KTest.fromfile(filename)
    # print('ktest filename : %r' % filename)
    # print('args       : %r' % b.args)
    print('num objects: %r' % len(b.objects))
    for i, (name, data) in enumerate(b.objects):
        str = trimZeros(data)
        print('object %4d: name: %r' % (i, name))
        print('object %4d: size: %r' % (i, len(data)))
        # if opts.writeInts and len(data) == 4:
        obj_data = struct.unpack('i', str)[0]
        print('object %4d: data: %r' %
              (i, obj_data))
        # else:
        # print('object %4d: data: %r' % (i, str))


def ktest_iterate():
    """ Iterate over all files ending with ktest in klee-last folder """

    directory = "klee-last"
    for filename in os.listdir(directory):
        if filename.endswith(".ktest"):
            file_list.append(os.path.join(directory, filename))
        else:
            continue

    file_list.sort()
    return file_list


def print_db_ktest():
    global file_list

    """ E - S For total task time
        R - L For total lock time
    """
    for filename in file_list:

        print("Testcase: %s" % filename)
        ktest_setdata(filename)
        """ Return all resources/tasks """
        c.execute('SELECT DISTINCT resource FROM events WHERE (file=?)',
                  (filename,))
        resources = c.fetchall()

        c.execute('SELECT DISTINCT job FROM events WHERE (file=?)',
                  (filename,))
        job = c.fetchall()
        print("Job: %s" % job[0])

        """ Go through the unique list of resources

            If they are called multiple times this is handled later
        """
        for res in resources:
            c.execute('''SELECT time FROM events WHERE
                      (file=? AND resource=?)''',
                      (str(filename), str(res[0]),))
            rows = c.fetchall()

            # Check if resource has been claimed more than than once
            claims = len(rows)

            # 2 is the expected value for 1 lock and 1 release
            x = 0
            while x < claims:
                print("Number of instructions taken for %s: %d"
                      % (res[0], rows[x+1][0] - rows[x][0]))
                x += 2

        print("\n")


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description=
                                     'Parse and print a profiling database.')
    parser.add_argument('filename', metavar='filename', type=open,
                        help='The database to be parsed')

    args = parser.parse_args()

    """ Save all ktest files into an array """
    file_list = ktest_iterate()

    print_db_ktest()