"""
Script to load example alerts data (as generated by alerts framework related
unittest) into a CouchDB database instance.

CouchDB URL and database name can be specified as command line arguments.

"""

import json
import logging
import sys
from contextlib import contextmanager
from argparse import ArgumentParser
from WMCore.Database.CMSCouch import CouchServer, Database


def _processCmdLineArgs(args):
    usage = \
        """usage: %prog options"""
    parser = ArgumentParser(usage=usage, add_help=False)
    _defineCmdLineOptions(parser)

    # opts - new processed options
    # args - remainder of the input array
    opts = parser.parse_args(args=args)
    for mandatory in ("input", "couchUrl", "database"):
        if not getattr(opts, mandatory, None):
            logging.error("Missing mandatory option ('%s')." % mandatory)
            parser.print_help()
            sys.exit(1)
    return opts.input, opts.couchUrl, opts.database


def _defineCmdLineOptions(parser):
    help = "Display this help"
    parser.add_argument("-h", "--help", help=help, action='help')
    help = "CouchDB server URL (mandatory)"
    parser.add_argument("-c", "--couchUrl", help=help)
    help = "CouchDB database name (mandatory)"
    parser.add_argument("-d", "--database", help=help)
    help = "Input data file (json format) (mandatory)"
    parser.add_argument("-i", "--input", help=help)


def _getDbConnection(couchUrl, dbName):
    """
    Check if the database exists, create if not.

    """
    couchServer = CouchServer(couchUrl)
    if not dbName in couchServer.listDatabases():
        logging.info("Database '%s' does not exits, creating it." % dbName)
        db = couchServer.createDatabase(dbName)  # returns Database
    else:
        logging.debug("Database '%s' exists." % dbName)
        db = Database(dbName, couchUrl)
    return couchServer, db


@contextmanager
def _handleFile(fileName):
    f = open(fileName, 'r')
    try:
        yield f
    finally:
        f.close()


def _processAndStoreFile(couchDb, inputFile):
    """
    Load input file.
    Each line corresponds to 1 alert (jsonized).
    Store one by one into the database.

    """
    decoder = json.decoder.JSONDecoder()
    with _handleFile(inputFile) as f:
        for alertStr in f:
            alert = decoder.decode(alertStr)
            couchDb.queue(alert)
    couchDb.commit()


def main():
    logging.basicConfig(level=logging.DEBUG)
    inputFile, couchUrl, dbName = _processCmdLineArgs(sys.argv)
    couchServer, couchDb = _getDbConnection(couchUrl, dbName)
    _processAndStoreFile(couchDb, inputFile)
    logging.info("Finished.")

    # helper stuff, checks (to remove later / if simon's happy):
    """
    allDocs = couchServer.get("/%s/_all_docs" % dbName)
    print "all documents, total_rows: %s" % allDocs["total_rows"]
    print allDocs

    # query a particular document (first) (need to know the id):
    id = allDocs["rows"][0]["id"]
    print "query document id: '%s'" % id
    db = Database(dbName, couchUrl)
    doc = db.document(id)
    print doc
    #couchServer.deleteDatabase(dbName)
    """


if __name__ == "__main__":
    main()
