#!/usr/bin/python
# -*- coding: utf-8 -*-
#
# Copyright 2009-2012 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 utilization efficency based on slot (machine) records.

# for floating point ops
from __future__ import division

# 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

# filter functions for the various states
def used_func(x): return x['State'] != "Unclaimed" and x['State'] != "Owner"
def unused_func(x): return x['State'] == "Unclaimed"
def owner_func(x): return x['State'] == "Owner"

# more accurate, using client-side filter (likely slower)
def single_query_reduce():
    total = list(raw['ads'].find({'MyType':'Machine'}))
    print "Total Slots:",len(total)
    used = filter(used_func,total)
    print "Used:",len(used)
    unused = filter(unused_func,total)
    print "Unused:",len(unused)
    owner = filter(owner_func,total)
    print "Owner:",len(owner)
    eff = (len(used)/(len(total)-len(owner)))*100
    print "Utilization: %.2f%%" % round(eff,2)

# less accurate but illustrates individual queries (faster)
def separate_queries():
    total = raw['ads'].find({'MyType':'Machine'}).distinct('Name')
    used = raw['ads'].find({'MyType':'Machine','State': {'$nin':['Unclaimed','Owner']}}).distinct('Name')
    unused = raw['ads'].find({'MyType':'Machine','State':'Unclaimed'}).distinct('Name')
    owner = raw['ads'].find({'MyType':'Machine','State':'Owner'}).distinct('Name')
    print "Total Slots:",len(total)
    print "Used:",len(used)
    print "Unused:",len(unused)
    print "Owner:",len(owner)
    eff = (len(used)/(len(total)-len(owner)))*100
    print "Utilization: %.2f%%" % round(eff,2)
    
def print_series(start,end):
    print "TIMESTAMP".ljust(30),"TOTAL".ljust(8),"USED".ljust(8),"UNUSED".ljust(8),"OWNER".ljust(8),"EFF %".ljust(8)
    report = stats['samples.machine'].find({'ts':{'$gte': parse(start)+UTC_DIFF, '$lt': parse(end)+UTC_DIFF}}).distinct('ts')
    for rec in report:
        total = stats['samples.machine'].find({'ts':rec}).distinct('n')
        used = stats['samples.machine'].find({'ts':rec,'st':{'$nin':['Unclaimed','Owner']}}).distinct('n')
        unused = stats['samples.machine'].find({'ts':rec,'st':'Unclaimed'}).distinct('n')
        owner = stats['samples.machine'].find({'ts':rec,'st':'Owner'}).distinct('n')
        eff = (len(used)/(len(total)-len(owner)))*100
        print str(rec).ljust(30), str(len(total)).ljust(8), str(len(used)).ljust(8), str(len(unused)).ljust(8), \
                str(len(owner)).ljust(8), "%.2f" % round(eff,2)

parser = OptionParser(description='Query Condor ODS for utilization efficency')
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('-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('-r','--remote', action="store_true",dest='remote',default=False, help='remotely query each slot state total (faster but less accurate)')
parser.add_option('-S','--series', action="store_true",dest="series",default=False, help='report pool utilization in start to end time range')
(options, args) =  parser.parse_args()

verbose = options.verbose

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

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

if options.remote:
    print "Querying remotely for each slot total..."
    separate_queries()
elif options.series:
    print_series(options.start, options.end)
else:
    print "Local calculation of slot totals..."
    single_query_reduce()
