#!/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, <wen.guan@cern.ch>, 2014
# - Cedric Serfon, <cedric.serfon@cern.ch>, 2015
#
import os
import sys

from rucio.client import Client
from rucio.common.config import config_get
from rucio.rse import rsemanager as rsemgr

OK, WARNING, CRITICAL, UNKNOWN = 0, 1, 2, 3


if __name__ == "__main__":

    cloud = sys.argv[1]

    retvalue = OK
    cloudRetValue = OK
    usedsize = 0
    freesize = 0

    try:
        proxy = config_get('nagios', 'proxy')
        os.environ["X509_USER_PROXY"] = proxy
    except Exception as e:
        print "Failed to get proxy from rucio.cfg"
        retvalue = WARNING

    c = Client()

    rses = c.list_rses('cloud=%s' % cloud)
    for rse in sorted(rses):
        rsename = rse['rse']
        rse_settings = rsemgr.get_rse_info(rsename)
        schemes = [p['scheme'] for p in rse_settings['protocols']]
        attrs = c.list_rse_attributes(rsename)

        if 'srm' not in schemes:
            print "The rse(%s) has no SRM protocol defined in AGIS" % (rsename)
            continue
        if not rse_settings['availability_read']:
            print "The rse(%s) blacklisted (not available for read)" % (rsename)
            continue

        for protocol in rse_settings['protocols']:
            if protocol['scheme'] == "srm":
                rse_settings['protocols'].remove(protocol)
                protocol['impl'] = "rucio.rse.protocols.gfal.Default"
                rse_settings['protocols'].append(protocol)
        try:
            gs, ret = rsemgr.get_space_usage(rse_settings, "srm")
            if gs:
                totalsize = long(ret["totalsize"])
                freesize = long(ret["unusedsize"])
                usedsize = totalsize - freesize
                retvalue = OK
            else:
                print "Failed to get rse(%s) space information: %s" % (rsename, str(ret))
                retvalue = WARNING
                if 'tier' in attrs and attrs['tier'] != '3' and 'type' in attrs and attrs['type'] != 'TEST':
                    cloudRetValue = WARNING
                else:
                    print "The rse(%s) is T3 or TEST => not setting overall failure" % (rsename)
        except Exception as e:
            print "Failed to get rse(%s) space information: %s" % (rsename, str(e))
            retvalue = WARNING
            if 'tier' in attrs and attrs['tier'] != '3' and 'type' in attrs and attrs['type'] != 'TEST':
                cloudRetValue = WARNING
            else:
                print "The rse(%s) is T3 or TEST => not setting overall failure" % (rsename)

        if retvalue == OK:
            print "Update RSE %s space usage (usedsize: %s, freesize: %s)" % (rsename, usedsize, freesize)
            c.set_rse_usage(rsename, "srm", usedsize, freesize)

    sys.exit(cloudRetValue)
