#!/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 dids, requests and replicas by hourly.
'''

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:
        lines = []
        session = get_session()
        utctime = datetime.datetime.utcnow()
        utctimeInt = int(time.mktime(utctime.timetuple()))
        timeStart = utctimeInt / 3600 * 3600 - 3600
        timeEnd = timeStart + 3600
        timezoneoffset = int((datetime.datetime.now() - datetime.datetime.utcnow()).seconds)

        # created datasets hourly
        sql = "select  /*+ index_ffs(dids DIDS_PK) */  count(1) from atlas_rucio.dids where scope='tests' and did_type='D' \
               and account='ddmusr01' and project='step14' and created_at>= to_timestamp('" + str(datetime.datetime.fromtimestamp(timeStart)) + "','YYYY-MM-dd HH24:MI:SS') \
               and created_at <to_timestamp('" + str(datetime.datetime.fromtimestamp(timeEnd)) + "','YYYY-MM-dd HH24:MI:SS')"
        result = session.execute(sql).fetchone()[0]
        monitor.record_gauge(stat='%s.hourly.created_datasets' % scope, value=result)
        print "hourly.created_datasets " + str(result)

        # created files hourly
        sql = "select  /*+ index_ffs(dids DIDS_PK) */  count(1) from atlas_rucio.dids where scope='tests' and did_type='F' \
              and account='ddmusr01' and created_at>= to_timestamp('" + str(datetime.datetime.fromtimestamp(timeStart)) + "','YYYY-MM-dd HH24:MI:SS') \
              and created_at <to_timestamp('" + str(datetime.datetime.fromtimestamp(timeEnd)) + "','YYYY-MM-dd HH24:MI:SS')"
        result = session.execute(sql).fetchone()[0]
        monitor.record_gauge(stat='%s.hourly.created_files' % scope, value=result)
        print "hourly.created_files " + str(result)

        # created replicas hourly
        sql = """SELECT
            CASE
                WHEN state = 'A' THEN 'replicas_available'
                WHEN state = 'U' THEN 'replicas_unavailable'
                WHEN state = 'C' THEN 'replicas_copying'
                WHEN state = 'B' THEN 'replicas_being_deleted'
                WHEN state = 'D' THEN 'replicas_bad'
                WHEN state = 'S' THEN 'replicas_source'
            ELSE state
        """
        sql += "END state_desc, num_rows FROM (select /*+ index_ffs((replicas REPLICAS_PK) */ state, count(*) num_rows FROM atlas_rucio.replicas \
                WHERE scope = 'tests' and created_at>= to_timestamp('" + str(datetime.datetime.fromtimestamp(timeStart)) + "','YYYY-MM-dd HH24:MI:SS') \
                and created_at <to_timestamp('" + str(datetime.datetime.fromtimestamp(timeEnd)) + "','YYYY-MM-dd HH24:MI:SS') \
                and rse_id != (select id from atlas_rucio.rses where rse = 'CERN-PROD-RUCIOTEST_DATADISK') GROUP BY state)"
        result = session.execute(sql).fetchall()
        for r in result:
            print 'hourly.created_%s ' % r[0], r[1]
            monitor.record_gauge(stat='%s.hourly.created_%s' % (scope, r[0]), value=r[1])

        # finished replicas hourly
        sql = "select  /*+ index_ffs(replicas REPLICAS_PK) */  count(1) from atlas_rucio.replicas where scope='tests' and state='A' \
               and updated_at>= to_timestamp('" + str(datetime.datetime.fromtimestamp(timeStart)) + "','YYYY-MM-dd HH24:MI:SS') \
               and updated_at <to_timestamp('" + str(datetime.datetime.fromtimestamp(timeEnd)) + "','YYYY-MM-dd HH24:MI:SS') \
               and rse_id != (select id from atlas_rucio.rses where rse = 'CERN-PROD-RUCIOTEST_DATADISK')"
        result = session.execute(sql).fetchone()[0]
        print 'hourly.finished_replicas %s' % result
        monitor.record_gauge(stat='%s.hourly.finished_replicas' % scope, value=result)

        # Total datasets hourly
        sql = "select  /*+ index_ffs(dids DIDS_PK) */  count(1) from atlas_rucio.dids where scope='tests' and did_type='D' and account='ddmusr01' and project='step14'"
        result = session.execute(sql).fetchone()[0]
        monitor.record_gauge(stat='%s.hourly.total_created_datasets' % scope, value=result)
        print "hourly.total_created_datasets " + str(result)

        # Total files hourly
        sql = "select  /*+ index_ffs(dids DIDS_PK) */  count(1) from atlas_rucio.dids where scope='tests' and did_type='F' and account='ddmusr01'"
        result = session.execute(sql).fetchone()[0]
        monitor.record_gauge(stat='%s.hourly.total_created_files' % scope, value=result)
        print "hourly.total_created_files " + str(result)

        # Total replicas hourly
        sql = """SELECT
            CASE
                WHEN state = 'A' THEN 'replicas_available'
                WHEN state = 'U' THEN 'replicas_unavailable'
                WHEN state = 'C' THEN 'replicas_copying'
                WHEN state = 'B' THEN 'replicas_being_deleted'
                WHEN state = 'D' THEN 'replicas_bad'
                WHEN state = 'S' THEN 'replicas_source'
            ELSE state
        """
        sql += "END state_desc, num_rows FROM (select /*+ index_ffs((replicas REPLICAS_PK) */ state, count(*) num_rows FROM atlas_rucio.replicas \
                WHERE scope = 'tests' and rse_id != (select id from atlas_rucio.rses where rse = 'CERN-PROD-RUCIOTEST_DATADISK') group by state)"
        result = session.execute(sql).fetchall()
        for r in result:
            print 'hourly.total_%s ' % r[0], r[1]
            monitor.record_gauge(stat='%s.hourly.total_%s' % (scope, r[0]), value=r[1])

        # Total requests hourly
        sql = """SELECT
            CASE
                WHEN state = 'Q' THEN 'requests_queued'
                WHEN state = 'S' THEN 'requests_submitted'
                WHEN state = 'F' THEN 'requests_failed'
                WHEN state = 'D' THEN 'requests_done'
                WHEN state = 'L' THEN 'requests_lost'
            ELSE state
            END state_desc, num_rows FROM (select /*+ index_ffs(REQUESTS REQUESTS_TYP_STA_CRE_IDX) */ state, count(*) num_rows
            FROM atlas_rucio.requests WHERE request_type='T' and scope = 'tests' GROUP BY state)
        """
        result = session.execute(sql).fetchall()
        for r in result:
            print 'hourly.total_%s ' % r[0], r[1]
            monitor.record_gauge(stat='%s.hourly.total_%s' % (scope, r[0]), value=r[1])

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