#!/usr/bin/env python
# Copyright European Organization for Nuclear Research (CERN) 2013
#
# 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
#
# Authors:
# - Wen Guan, <wguan@cern.ch>, 2014

'''
Probe to check requests.
'''

import datetime
import sys
import time

from rucio.core import monitor
from rucio.db.sqla.session import get_session

# Exit statuses
OK, WARNING, CRITICAL, UNKNOWN = 0, 1, 2, 3


if __name__ == "__main__":

    scope = "stresstest"

    try:
        session = get_session()
        utctime = datetime.datetime.utcnow()
        utctimeInt = int(time.mktime(utctime.timetuple()))
        timeEnd = utctimeInt
        timezoneoffset = int((datetime.datetime.now() - datetime.datetime.utcnow()).seconds)

        # check requests to different sites
        sql = "select /*+ index_ffs(REQUESTS REQUESTS_TYP_STA_CRE_IDX) */ atlas_rucio.id2rse(dest_rse_id) rse, state, count(*) \
               from atlas_rucio.requests where request_type='T'  group by dest_rse_id, state"
        result = session.execute(sql).fetchall()
        for rse, state, num in result:
            # print "requests.%s.%s %s" % (state, rse, num)
            monitor.record_gauge(stat='%s.requests.%s.%s' % (scope, state, rse), value=num)

        # check left requests which are 12 hours old to different sites
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=12)
        sql = "select /*+ index_ffs(REQUESTS REQUESTS_TYP_STA_CRE_IDX) */ atlas_rucio.id2rse(dest_rse_id) rse, state, count(*) \
              from atlas_rucio.requests where request_type='T' and created_at<=to_timestamp('" + str(timeLimit) + "','YYYY-MM-dd HH24:MI:SS.FF') group by dest_rse_id, state"
        result = session.execute(sql).fetchall()
        for rse, state, num in result:
            # print 'requests.12Hours_old.%s.%s %s' % (state, rse, num)
            monitor.record_gauge(stat='%s.requests.12Hours_old.%s.%s' % (scope, state, rse), value=num)

        # check left requests which are 24 hours old to different sites
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=24)
        sql = "select /*+ index_ffs(REQUESTS REQUESTS_TYP_STA_CRE_IDX) */ atlas_rucio.id2rse(dest_rse_id) rse, state, count(*) \
              from atlas_rucio.requests where request_type='T' and created_at<=to_timestamp('" + str(timeLimit) + "','YYYY-MM-dd HH24:MI:SS.FF') group by dest_rse_id, state"
        result = session.execute(sql).fetchall()
        for rse, state, num in result:
            # print 'requests.24Hours_old.%s.%s %s' % (state, rse, num)
            monitor.record_gauge(stat='%s.requests.24Hours_old.%s.%s' % (scope, state, rse), value=num)

        # check left requests which are 48 hours old to different sites
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=48)
        sql = "select /*+ index_ffs(REQUESTS REQUESTS_TYP_STA_CRE_IDX) */ atlas_rucio.id2rse(dest_rse_id) rse, state, count(*) from atlas_rucio.requests \
               where request_type='T' and created_at<=to_timestamp('" + str(timeLimit) + "','YYYY-MM-dd HH24:MI:SS.FF') group by dest_rse_id, state"
        result = session.execute(sql).fetchall()
        for rse, state, num in result:
            # print 'requests.48Hours_old.%s.%s %s' % (state, rse, num)
            monitor.record_gauge(stat='%s.requests.48Hours_old.%s.%s' % (scope, state, rse), value=num)

    except:
        sys.exit(UNKNOWN)
    finally:
        session.remove()
    sys.exit(OK)
