#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright 2009-2011 Red Hat, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# This utility script is designed to query the mongodb database
# for sampled records of Condor machine slots and submitter job
# totals.

# uses pymongo - http://pypi.python.org/pypi/pymongo/
import pymongo
from datetime import timedelta, datetime
from sys import exit, argv
import time, pwd, os
from optparse import OptionParser
from dateutil import *
from dateutil.parser import *

# NOTE: the 1.9 pymongo driver does implicit localtime conversion when
# dates are used in a query so we must compensate, eventhough the dates
# ARE already stored in localtime in 1.6.4 mongodb
# mongodb 1.7+ uses ISODate to address this
UTC_DIFF = datetime.utcnow() - datetime.now()
DEFAULT_START_DT = str(datetime.utcnow()-UTC_DIFF-timedelta(minutes=15))
DEFAULT_END_DT = str(datetime.utcnow()-UTC_DIFF)

verbose = False

def print_user(user,start,end):
    print "SUBMITTER".ljust(40),"TIMESTAMP".ljust(30),"R".ljust(5),"H".ljust(5),"I"
    for item in db['samples.submitter'].find({"sn":{'$regex':'^'+user}, 'ts':{'$gte': parse(start)+UTC_DIFF, '$lt': parse(end)+UTC_DIFF}}):
        print item['sn'].ljust(40),str(item['ts']-UTC_DIFF).ljust(30), str(item['jr']).ljust(5),str(item['jh']).ljust(5),str(item['ji'])

def print_resource(resource,start,end):
    print "MACHINE".ljust(40),"TIMESTAMP".ljust(30),"ARCH/OS".ljust(15),"IDLE".ljust(8),"LOAD".ljust(5),"STATE"
    for item in db['samples.machine'].find({"n":{'$regex':resource}, 'ts':{'$gte': parse(start)+UTC_DIFF, '$lt': parse(end)+UTC_DIFF}}):
        print item['n'].ljust(40),str(item['ts']-UTC_DIFF).ljust(30), str("%s/%s" % (item['ar'],item['os'])).ljust(15),str(item['ki']).ljust(8),\
        str(item['la'])[:5].ljust(5),item['st']

def print_users(start,end):
    print "SUBMITTER"
    for user in db['samples.submitter'].find({'ts':{'$gte': parse(start)+UTC_DIFF, '$lt': parse(end)+UTC_DIFF}}).distinct('sn'):
        print user

def print_usergroups(start,end):
    # NOTE: after revisiting, mongodb group() seems significantly slower than find()/distinct() below
    # ie., db['samples.submitter'].group({key: {'sn': true,'mn':true }, cond: {}, initial: {}, reduce: function() {}});
    #usergroups = db['samples.submitter'].group(['sn','mn'], None, {}, 'function() {}')
    #for ug in usergroups:
        #print '%s/%s' % (ug['sn'],ug['mn'])
    print "SUBMITTER / MACHINE"
    mnlist = db['samples.submitter'].find({'ts':{'$gte': parse(start)+UTC_DIFF, '$lt': parse(end)+UTC_DIFF}}).distinct('ma')
    for mn in mnlist:
        snlist = db['samples.submitter'].find({'ma': mn, 'ts':{'$gte': parse(start)+UTC_DIFF, '$lt': parse(end)+UTC_DIFF}}).distinct('sn')
        for sn in snlist:
            print '%s / %s' % (sn,mn)

def print_resourcelist(start,end):
    print "MACHINE"
    for item in db['samples.machine'].find({'ts':{'$gte': parse(start)+UTC_DIFF, '$lt': parse(end)+UTC_DIFF}}).distinct('n'):
        print item

parser = OptionParser(description='Query HTCondor ODS for statistics')
parser.add_option('-v','--verbose', action="store_true",default=False, help='enable logging')
parser.add_option('-s','--server', action="store", dest='server',
                    default='localhost',
                    help='mongodb database server location: e.g., somehost, localhost:2011')
parser.add_option('--u','--user', dest="user", help='stats for a single submitter: user,timestamp,running,held,idle')
parser.add_option('--r','--resource', dest="resource", help='stats for a single resource: slot,timestamp,keyboard idle,load average,status')
parser.add_option('--f','--from', dest="start", help='records from datetime in ISO8601 format e.g., \'2011-09-29 12:03\'', default=DEFAULT_START_DT)
parser.add_option('--t','--to', dest="end", help='records to datetime in ISO8601 format e.g., \'2011-09-30T17:16\'',default=DEFAULT_END_DT)
parser.add_option('--ul','--userlist', action="store_true",dest="userlist", default=False, help='list all submitters')
parser.add_option('--ugl','--usergrouplist', action="store_true",dest='usergrouplist',default=False, help='list all submitter groups')
parser.add_option('--rl','--resourcelist', action="store_true",dest='resourcelist',default=False, help='list all resources')
(options, args) =  parser.parse_args()

verbose = options.verbose

try:
    connection = pymongo.Connection(options.server)
    db = connection.condor_stats
except Exception, e:
    print e
    exit(1)

if verbose:
    print 'from:\t', options.start
    print 'to:\t', options.end

if options.user:
    print_user(options.user,options.start,options.end)
elif options.resource:
    print_resource(options.resource,options.start,options.end)
elif options.userlist:
    print_users(options.start,options.end)
elif options.usergrouplist:
    print_usergroups(options.start,options.end)
elif options.resourcelist:
    print_resourcelist(options.start,options.end)
else:
    parser.print_help()
