#!/usr/bin/env python
"""
Tag start/end of experiment by making an entry in MongoDB.
"""
# Stdlib
import sys
import optparse
import pprint
import time
# Third party
import pymongo

DEFAULT_DB = "periscope"
DEFAULT_COLL = "experiments"

def run(tag=None, mode=None, db=None, coll=None, meta=None):
    ts = time.time()
    conn = pymongo.Connection()
    coll = conn[db][coll]
    coll.ensure_index('tag', unique=True)
    spec = {'tag':tag}
    if mode == 'start':
        record = {
            'start' : ts,
            'end' : -1,
            'tag' : tag,
            'meta' : meta }
        coll.update(spec, record, upsert=True)
    elif mode == 'end':
        coll.update(spec, {'$set' : {'end' : ts}})
    elif mode == 'del':
        coll.remove(spec)
    elif mode == 'show':
        records = list(coll.find(spec))
        if not records:
            print("NOT-FOUND: {tag}".format(tag=tag))
        else:
            for record in records:
                d = record.copy()
                del d['_id']
                pprint.pprint(d)
    else:
        raise ValueError("bad mode {0}".format(mode))

def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]
    desc = ' '.join(__doc__.split())
    MODES = ("start", "end", "del", "show")
    modes_arg = "|".join(MODES)
    op = optparse.OptionParser(usage="%prog [options] <tag> {0} "
                               "[key1=desc1 key2=desc2 ..]".format(modes_arg),
                               description=desc)
    op.add_option("-c", "--collection", dest="coll", metavar="NAME",
                  default=DEFAULT_COLL,
                  help="Mongo collection (default=%default)")
    op.add_option("-d", "--db", dest="db", metavar="NAME", default=DEFAULT_DB,
                  help="Mongo database (default=%default)")
    op.add_option("-u", "--unique", dest="unique", action="store_true",
                  help="For mode 'start', add some hex digits to "
                  "provided 'tag' to make it unique")
    options, args = op.parse_args(argv)
    if len(args) < 2:
        op.error("First two arguments are required")
    tag = args[0]
    mode = args[1].lower()
    
    if mode not in MODES:
        op.error("Second argument must be {0}".format(modes_arg))
    if mode != 'start' and options.unique:
        op.error("-u/--unique only valid for mode 'start'")
    kw = { }
    for kvp in args[2:]:
        if not '=' in kvp:
            op.error("Bad format for key=value pair: {0}".format(kvp))
        key, value = kvp.split("=", 1)
        kw[key] = value
    if options.unique:
        import uuid
        tag = "{tag}_{foo}".format(tag=tag, foo=str(uuid.uuid1()))
    run(tag=tag, mode=mode, db=options.db, coll=options.coll, meta=kw)
    meta_str = ' '.join([k+'='+v for k,v in kw.items()])
    print("{mode} {tag} {meta}".format(
        mode=mode, tag=tag, meta=meta_str))
    return 0

if __name__ == '__main__':
    sys.exit(main())
