#!/usr/bin/env python
# Ok so I know this is basically a bs metric, but it's still fun to see what the numbers
# are.  But be warned, don't take these numbers too seriously.
import sys
import shutil
import time
import os
import tempfile
import optparse


def main():
    parser = optparse.OptionParser()
    parser.add_option('-n', '--num-transactions', default=1000000, type=int)
    parser.add_option('-c', '--chunk-size', default=10000, type=int,
                      help="For the read chunked tests, this will "
                      "set how many elements to iterate over at a time")
    parser.add_option('-r', '--repeat', default=10, type=int,
                      help="For the read chunked tests, this will "
                      "specify how many times to iterate over the chunks "
                      "before moving on to the next chunk.")
    parser.add_option('-t', '--skip-read-test', action="store_true",
                      default=False, help="Skip the sequential read test "
                      "(useful if you just want to benchmark writes)")
    parser.add_option('-s', '--skip-read-chunk', action="store_true",
                      default=False, help="Skip the read chunk test "
                      "(it can take a while to run).")
    parser.add_option('-d', '--dbm', default='semidbm')
    opts, args = parser.parse_args()

    tempdir = tempfile.mkdtemp(prefix='tps')
    dbname = os.path.join(tempdir, 'tps.db')
    try:
        dbm_module = __import__(opts.dbm, fromlist=[opts.dbm])
    except ImportError:
        sys.stderr.write("Can't import dbm: %s\n" % opts.dbm)
        sys.exit(1)
    db = dbm_module.open(dbname, 'c')

    num_transactions = opts.num_transactions
    groups_of = opts.chunk_size
    repeat = opts.repeat

    start = time.time()
    for i in xrange(num_transactions):
        db[str(i)] = str(i)
    end = time.time()
    print "Write ",
    print "Total: %.5f, tps: %.2f" % (end - start,
                                      float(num_transactions) / (end - start))
    if not opts.skip_read_test:
        db.close()
        db = dbm_module.open(dbname, 'r')
        start = time.time()
        for i in xrange(num_transactions):
            db[str(i)]
        end = time.time()

        print "Read ",
        print "Total: %.5f, tps: %.2f" % (end - start,
                                          float(num_transactions) / (end - start))


    if not opts.skip_read_chunk:
        count = 0
        start = time.time()
        for i in xrange(0, num_transactions, groups_of):
            for j in xrange(groups_of):
                for k in xrange(repeat):
                    count += 1
                    db[str(i + j)]
        end = time.time()
        print "Read (grouped)",
        print "count:", count
        print "Total: %.5f, tps: %.2f" % (end - start,
                                          float(count) / (end - start))
    db.close()
    shutil.rmtree(tempdir)


if __name__ == '__main__':
    main()