#!/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
#
import argparse
import os
import sys

from rucio.client import Client
from rucio.rse import rsemanager as rsemgr

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


if __name__ == "__main__":
    usage = """
usage: %(prog)s [args]

"""
    oparser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]), add_help=True)
    oparser.add_argument('--RSE', dest="rsename", default=None, help="The RSE name")

    if len(sys.argv) == 1:
        oparser.print_help()
        sys.exit(WARNING)
    else:
        args = oparser.parse_args(sys.argv[1:])
        rsename = args.rsename

    c = Client()

    retvalue = OK
    usedsize = 0
    freesize = 0

    rse_settings = rsemgr.get_rse_info(rsename)
    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
        else:
            print "Failed to get rse(%s) space information: %s" % (rsename, str(ret))
            retvalue = CRITICAL
    except Exception as e:
        print "Failed to get rse(%s) space information: %s" % (rsename, str(e))
        retvalue = CRITICAL

    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(retvalue)
