#!/usr/bin/env python
"""Very simple script for profiling various dbms.

The point of this script is to give a rough
estimate for how semidbm does compared to other
dbms.  You can run this script with no args or
specify the dbms you want to benchmark using
the --dbm arg.

"""
import os
import sys
import stat
import json
import shutil
import optparse
import time
import string
import tempfile
import random
import traceback

try:
    _range = xrange
except NameError:
    _range = range

random.seed(100)


_potential_dbms = ['dbhash', 'dbm', 'gdbm', 'dumbdbm', 'semidbm']

ADAPTER_DIR = os.path.join(os.path.dirname(__file__), 'adapters')
sys.path.append(ADAPTER_DIR)
out = sys.stdout.write


def _rand_key(key_length, chars=string.printable):
    return bytes(bytearray(''.join(random.choice(chars) for i in
                                   _range(key_length))))


def set_dbms(dbms):
    dbms_found = []
    for potential in dbms:
        try:
            d = __import__(potential, fromlist=[potential])
            dbms_found.append(d)
        except ImportError as e:
            sys.stderr.write("Could not import %s: %s\n" % (potential, e))
            continue
    return dbms_found


class Options(object):
    num_keys = 1000000
    key_size_bytes = 16
    value_size_bytes = 100

    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

    def print_options(self):
        stats = ("    num_keys  : %(num_keys)s\n"
                 "    key_size  : %(key_size_bytes)s\n"
                 "    value_size: %(value_size_bytes)s" % self.__dict__)
        return stats

    @property
    def key_format(self):
        return '%0' + str(self.key_size_bytes) + 'd'


class StatsReporter(object):
    def __init__(self, name, total_time, total_bytes, total_ops):
        self._name = name
        self._total_time = total_time
        self._total_bytes = total_bytes
        self._total_ops = total_ops

    def micros_per_op(self):
        # Leveldb uses this, so it's useful to compare.
        total_micros = self._total_time * 1e6
        return total_micros / self._total_ops

    def ops_per_second(self):
        return self._total_ops / float(self._total_time)

    def megabytes_per_second(self):
        return self._total_bytes / (1024.0 * 1024) / self._total_time

    def print_report(self):
        out("%-20s:" % self._name)
        out(" time: %9.3f,   micros/ops: %9.3f,   ops/s: %10.3f,  "
              "MB/s: %10.3f\n" % (self._total_time, self.micros_per_op(),
                                self.ops_per_second(),
                                self.megabytes_per_second()))

    @property
    def name(self):
        return self._name


class Benchmarks(object):
    def __init__(self, options, tmpdir):
        self.options = options
        self.tmpdir = tmpdir
        self.random_values = self._generate_random_string(1024 * 1024)

    def _generate_random_string(self, string_size):
        out("Generating random data.\n")
        c = chr
        rand = random.randint
        r = bytes(bytearray([rand(0, 255) for i in _range(string_size)]))
        return r

    def run(self, dbm):
        print("Benchmarking:", dbm)
        print(self.options.print_options())
        all_reports = []
        try:
            for name in ['fill_random', 'fill_sequential', 'read_cold',
                         'read_sequential', 'read_hot', 'read_random',
                         'delete_sequential']:
                method = getattr(self, name)
                report = method(dbm)
                report.print_report()
                all_reports.append(report)
        finally:
            self.delete_dbm()
            print
        return all_reports

    def fill_random(self, dbm):
        db = self._load_dbm(dbm)
        random_values = self.random_values
        maxlen = len(random_values)
        position = 0
        value_size = self.options.value_size_bytes
        key_size = self.options.key_size_bytes
        num_keys = self.options.num_keys
        indices = [_rand_key(key_size) for i in _range(num_keys)]

        t = time.time
        out = sys.stdout.write
        flush = sys.stdout.flush
        start = t()
        for i in _range(num_keys):
            db[indices[i]] = random_values[position:position+value_size]
            position += value_size
            if position + value_size > maxlen:
                position = 0
                out("(%s/%s)\r" % (i, num_keys))
                flush()
        total = t() - start
        self._close_db(db)
        self.delete_dbm()
        return StatsReporter(
            'fill_random', total,
            (value_size * num_keys) + (self.options.key_size_bytes * num_keys),
            num_keys)

    def fill_sequential(self, dbm):
        db = self._load_dbm(dbm)
        key_format = self.options.key_format
        random_values = self.random_values
        maxlen = len(random_values)
        position = 0
        value_size = self.options.value_size_bytes
        num_keys = self.options.num_keys
        indices = [(key_format % i).encode('utf-8') for i in _range(num_keys)]

        t = time.time
        out = sys.stdout.write
        flush = sys.stdout.flush
        start = t()
        for i in _range(num_keys):
            db[indices[i]] = random_values[position:position+value_size]
            position += value_size
            if position + value_size > maxlen:
                position = 0
                out("(%s/%s)\r" % (i, num_keys))
                flush()
        total = t() - start
        self._close_db(db)
        return StatsReporter(
            'fill_sequential', total,
            (value_size * num_keys) + (self.options.key_size_bytes * num_keys),
            num_keys)

    def read_sequential(self, dbm, name='read_sequential'):
        # Assumes fill_sequential has been called.
        db = self._load_dbm(dbm, 'r')
        key_format = self.options.key_format
        num_keys = self.options.num_keys

        indices = [(key_format % i).encode('utf-8') for i in _range(num_keys)]
        t = time.time
        start = t()
        for i in _range(num_keys):
            db[indices[i]]
        total = t() - start
        self._close_db(db)
        total_bytes = (self.options.key_size_bytes * num_keys +
                       self.options.value_size_bytes * num_keys)
        return StatsReporter(name, total, total_bytes, num_keys)

    def read_cold(self, dbm):
        # read_cold is intended to be called before read_sequential or any
        # other reads to test the performance of a "cold" read.
        return self.read_sequential(dbm, name='read_cold')

    def read_hot(self, dbm):
        # Assumes fill_sequential has been called.
        # Read from 1% of the database self.options.num_keys times.
        # This should test the effectiveness of any caching being used.
        num_keys = self.options.num_keys
        unique_keys = int(num_keys * 0.01)
        indices = [(self.options.key_format % i).encode('utf-8')
                   for i in random.sample(_range(num_keys), unique_keys)]
        indices = indices * (int(num_keys / unique_keys))
        db = self._load_dbm(dbm, 'r')
        t = time.time
        start = t()
        for i in _range(num_keys):
            db[indices[i]]
        total = t() - start
        self._close_db(db)
        total_bytes = (self.options.key_size_bytes * num_keys +
                       self.options.value_size_bytes * num_keys)
        return StatsReporter('read_hot', total, total_bytes,
                             num_keys)

    def read_random(self, dbm):
        # This doesn't matter to semidbm because the keys
        # aren't ordered, but other dbms might be impacted.
        num_keys = self.options.num_keys
        key_format = self.options.key_format
        indices = [(key_format % i).encode('utf-8') for i in range(num_keys)]
        random.shuffle(indices)
        db = self._load_dbm(dbm, 'r')
        t = time.time
        start = t()
        for i in _range(num_keys):
            db[indices[i]]
        total = t() - start
        self._close_db(db)
        total_bytes = (self.options.key_size_bytes * num_keys +
                       self.options.value_size_bytes * num_keys)
        return StatsReporter('read_random', total, total_bytes,
                             num_keys)

    def delete_sequential(self, dbm):
        # Assumes fill_sequential has been called.
        db = self._load_dbm(dbm, 'c')
        key_format = self.options.key_format
        num_keys = self.options.num_keys

        indices = [(key_format % i).encode('utf-8') for i in _range(num_keys)]
        t = time.time
        start = t()
        for i in _range(num_keys):
            del db[indices[i]]
        total = t() - start
        self._close_db(db)
        total_bytes = (self.options.key_size_bytes * num_keys +
                       self.options.value_size_bytes * num_keys)
        return StatsReporter('delete_sequential', total, total_bytes, num_keys)

    def delete_dbm(self):
        # Just wipe out everything under tmpdir.
        self._rmtree(self.tmpdir)

    def _rmtree(self, tmpdir):
        # Delete everything under tmpdir but don't actually
        # delete tmpdir itself.
        for path in os.listdir(tmpdir):
            full_path = os.path.join(tmpdir, path)
            mode = os.lstat(full_path).st_mode
            if stat.S_ISDIR(mode):
                shutil.rmtree(full_path)
            else:
                os.remove(full_path)

    def _load_dbm(self, dbm, flags='c'):
        db = dbm.open(os.path.join(self.tmpdir, 'db'), flags)
        return db

    def _close_db(self, db):
        # If the db has a close() method call it.  Basically a hack
        # so we can benchmark a normal python dict.
        if hasattr(db, 'close'):
            db.close()


def generate_report(filename, options, reports):
    """Create a json report grouped by benchmarks rather than by dbm.

    Since this is going to be used to autogenerate the
    charts/tables, a comparison across dbms for a given benchmark
    is more useful.  The output should look like::

        {num_keys: 100, key_size_bytes: 16, value_size_bytes: 1000,
         dbms: ['semidbm', 'gdbm'],
         benchmarks:
             [['fill_sequential', [
                {total_time: 100, micros_per_op: 1,
                 ops_per_second: 123, mb_per_second: 100}]],
              ...
             ]
        }


    """
    # Generating a report requires python >= 2.7.
    from collections import OrderedDict
    output = {
        'num_keys': options.num_keys,
        'key_size_bytes': options.key_size_bytes,
        'value_size_bytes': options.value_size_bytes
    }
    by_benchmarks = OrderedDict()
    dbms = []
    for dbm, benchmarks in reports:
        dbms.append(dbm)
        for benchmark in benchmarks:
            by_benchmarks.setdefault(benchmark.name, []).append({
                'total_time': benchmark.total_time(),
                'micros_per_op': benchmark.micros_per_op(),
                'ops_per_second': benchmark.ops_per_second(),
                'megabytes_per_second': benchmark.megabytes_per_second(),
            })
    output['dbms'] = dbms
    output['benchmarks'] = by_benchmarks
    json.dump(output, open(filename, 'w'), indent=4)


def main():
    parser = optparse.OptionParser()
    parser.add_option('-d', '--dbm', dest='dbms', action='append')
    # These are the same defaults as the leveldb benchmark,
    # which this scripts is based off of.
    parser.add_option('-n', '--num-keys', default=1000000, type=int)
    parser.add_option('-k', '--key-size-bytes', default=16, type=int)
    parser.add_option('-s', '--value-size-bytes', default=100, type=int)
    parser.add_option('-r', '--report', help="Generate a summary report "
                      "in json to specified location.")
    opts, args = parser.parse_args()


    dbm_names = opts.__dict__.pop('dbms') or _potential_dbms
    dbms = set_dbms(dbm_names)
    if not dbms:
        sys.stderr.write("List of dbms is empty.\n")
        sys.exit(1)
    options = Options(**opts.__dict__)
    tmpdir = tempfile.mkdtemp(prefix='dbmprofile')
    benchmarks = Benchmarks(options, tmpdir)
    all_reports = []
    try:
        for dbm in dbms:
            try:
                all_reports.append((dbm.__name__, benchmarks.run(dbm)))
            except Exception as e:
                traceback.print_exc()
                sys.stderr.write(
                    "ERROR: exception caught when benchmarking %s: %s\n" %
                    (dbm, e))
    finally:
        shutil.rmtree(tmpdir)
    if opts.report:
        generate_report(opts.report, options, all_reports)


if __name__ == '__main__':
    main()