#!/usr/bin/env python

import sys, datetime, os

def run():
    uri = 'postgres://scott:tiger@127.0.0.1/test'

    test = ZooMark()

    test.setup(uri, ITERATIONS)
    try:
        total = datetime.timedelta(0) 
        for method in [x for x in dir(test) if x.startswith("step_")]:
            startTime = datetime.datetime.now()
            meth = getattr(test, method)
            try:
                if profile or callcount:
                    funccount = profile_callable(meth, profile)
                else:
                    meth()
            except:
                import traceback
                traceback.print_exc()
            testtime = datetime.datetime.now() - startTime
            if 'funccount' in locals():
                print "%s/%s/%s : %d" % (title, ITERATIONS, method, funccount)
            else:
                print "%s/%s/%s : %s" % (title, ITERATIONS, method, testtime)
            
            # dont include step 7 in total since Storm tests dont implement
            if not x.startswith("step_7"):
                total += testtime
    finally:
        if not callcount:
            print "%s/%s/%s : %s" % (title, ITERATIONS, "total_tests_1_thru_6", total)
        test.teardown()

def profile_callable(fn, print_=True):
    import time, hotshot, hotshot.stats

    target = fn.__name__
    filename = "%s.prof" % target

    prof = hotshot.Profile(filename)
    began = time.time()
    prof.start()
    try:
        fn()
    finally:
        prof.stop()
        ended = time.time()
        prof.close()

    if print_:
        print "Profiled target '%s', wall time: %.2f seconds" % (target, ended - began)
        print "Profile report for target '%s' (%s)" % (target, filename)

    stats = hotshot.stats.load(filename)
    if print_:
        stats.sort_stats('cumulative')
        stats.print_stats()

    try:
        return stats.total_calls
    finally:
        os.unlink(filename)

if __name__ == '__main__':
    profile = False
    callcount = False
    module = 'sqla'
    ITERATIONS = 100
    title = None

    args = list(sys.argv)
    while args:
        arg = args.pop(0)
        if arg == '--profile':
            profile = True
        elif arg == '--module':
            module = args.pop(0)
        elif arg == '--callcount':
            callcount = True
        elif arg == '--title':
            title = args.pop(0)
        elif not args:
            ITERATIONS = int(arg)

    if not title:
        title = module 

    execfile("./benchmark_%s.py" % module)

    run()