#!/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 rules.
'''

import datetime
import sys
import time

from rucio.core import monitor
from rucio.db.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 rules
        sql = """SELECT
            CASE
                WHEN state = 'R' THEN 'rules_replicating'
                WHEN state = 'O' THEN 'rules_ok'
                WHEN state = 'S' THEN 'rules_stuck'
                WHEN state = 'U' THEN 'rules_suspend'

            ELSE state
        END state_desc, num_rows FROM (select /*+ index_ffs((rules RULES_PK) */ state, count(*) num_rows
        from atlas_rucio.rules where scope='tests' and account='ddmadmin' and rse_expression!='CERN-PROD-RUCIOTEST_DATADISK' group by state)
        """

        result = session.execute(sql).fetchall()
        for state, num in result:
            print 'rules.%s %s' % (state, num)
            monitor.record_gauge(stat='%s.rules.%s' % (scope, state), value=num)

        # check stuck rules
        sql = "select sum(locks_stuck_cnt)  from atlas_rucio.rules where scope='tests' and account='ddmadmin' and rse_expression!='CERN-PROD-RUCIOTEST_DATADISK' and state='S' "
        result = session.execute(sql).fetchone()[0] or 0
        print 'rules.sum_locks_stuck_cnt %s' % (result)
        monitor.record_gauge(stat='%s.rules.sum_locks_stuck_cnt' % scope, value=result)

        # check stuck rules older than 24 hours
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=24)
        sql = "select sum(locks_stuck_cnt)  from atlas_rucio.rules where scope='tests' and account='ddmadmin' \
               and rse_expression!='CERN-PROD-RUCIOTEST_DATADISK' and state='S' and created_at <= to_timestamp('"+str(timeLimit)+"','YYYY-MM-dd HH24:MI:SS.FF')"
        result = session.execute(sql).fetchone()[0] or 0
        print 'rules.created_24hours_ago.sum_locks_stuck_cnt %s' % (result)
        monitor.record_gauge(stat='%s.rules.created_24hours_ago.sum_locks_stuck_cnt' % scope, value=result)

        # check stuck rules older than 1 week
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=168)
        sql = "select sum(locks_stuck_cnt)  from atlas_rucio.rules where scope='tests' and account='ddmadmin' \
              and rse_expression!='CERN-PROD-RUCIOTEST_DATADISK' and state='S' and created_at <= to_timestamp('"+str(timeLimit)+"','YYYY-MM-dd HH24:MI:SS.FF')"
        result = session.execute(sql).fetchone()[0] or 0
        print 'rules.created_1week_ago.sum_locks_stuck_cnt %s' % (result)
        monitor.record_gauge(stat='%s.rules.created_1week_ago.sum_locks_stuck_cnt' % scope, value=result)

        # check replicating rules
        sql = "select sum(locks_replicating_cnt)  from atlas_rucio.rules where scope='tests' and account='ddmadmin' and rse_expression!='CERN-PROD-RUCIOTEST_DATADISK' and state in ('S','R') "
        result = session.execute(sql).fetchone()[0] or 0
        print 'rules.sum_locks_replicating_cnt %s' % (result)
        monitor.record_gauge(stat='%s.rules.sum_locks_replicating_cnt' % scope, value=result)

        # check replicating rules older than 24 hours
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=24)
        sql = "select sum(locks_replicating_cnt)  from atlas_rucio.rules where scope='tests' and account='ddmadmin' \
               and rse_expression!='CERN-PROD-RUCIOTEST_DATADISK' and state in ('S','R') and created_at <= to_timestamp('"+str(timeLimit)+"','YYYY-MM-dd HH24:MI:SS.FF')"
        result = session.execute(sql).fetchone()[0] or 0
        print 'rules.created_24hours_ago.sum_locks_replicating_cnt %s' % (result)
        monitor.record_gauge(stat='%s.rules.created_24hours_ago.sum_locks_replicating_cnt' % scope, value=result)

        # check replicating rules older than 1 week
        timeLimit = datetime.datetime.utcnow() - datetime.timedelta(hours=168)
        sql = "select sum(locks_replicating_cnt)  from atlas_rucio.rules where scope='tests' and account='ddmadmin' \
              and rse_expression!='CERN-PROD-RUCIOTEST_DATADISK' and state in ('S','R') and created_at <= to_timestamp('"+str(timeLimit)+"','YYYY-MM-dd HH24:MI:SS.FF')"
        result = session.execute(sql).fetchone()[0] or 0
        print 'rules.created_1week_ago.sum_locks_replicating_cnt %s' % (result)
        monitor.record_gauge(stat='%s.rules.created_1week_ago.sum_locks_replicating_cnt' % scope, value=result)

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