#!/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
 - Sylvain Blunier, <sylvain.blunier@cern.ch>, 2016
 - Tomas Javurek, <Tomas.Javurek@cern.ch>, 2016
 - Cedric Serfon, <cedric.serfon@cern.ch>, 2015-2017
 - Dimitrios Christidis, <dimitrios.christidis@cern.ch>, 2019
"""

import json
import os
import sys

import gfal2
import requests

from rucio.client import Client
from rucio.common.config import config_get
from rucio.common.exception import RSEProtocolNotSupported
from rucio.rse import rsemanager as rsemgr
from rucio.api.rse import list_rse_attributes, get_rse

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

gfal2.set_verbose(gfal2.verbose_level.normal)


def consume_atlas_json(data, key):
    totalsize = int(data[key]['total_space'])
    used = int(data[key]['used_space'])
    unusedsize = totalsize - used
    return totalsize, unusedsize


def consume_wlcg_json(data, key):
    index = {share['name']: share for share in data['storageservice']['storageshares'] if 'name' in share}
    totalsize = index[key]['totalsize']
    used = index[key]['usedsize']
    unusedsize = totalsize - used
    return totalsize, unusedsize


def get_space_usage(rse_name):
    """
    Get RSE space usage information.

    :returns: a list with dict containing 'totalsize' and 'unusedsize'

    """
    dest = '/tmp/rucio-gsiftp-site-size_' + rse_name
    space_usage_url = ''
    # url of space usage json, woud be nicer to have it in rse_settings
    agis = requests.get('http://atlas-agis-api.cern.ch/request/ddmendpoint/query/list/?json').json()
    agis_token = ''
    for res in agis:
        if rse_name == res['name']:
            agis_token = res['token']
            space_usage_url = res['space_usage_url']
    print space_usage_url
    if os.path.exists(dest):
        os.remove(dest)
    ctx = gfal2.creat_context()
    ctx.set_opt_string_list("SRM PLUGIN", "TURL_PROTOCOLS", ["gsiftp", "rfio", "gsidcap", "dcap", "kdcap"])
    params = ctx.transfer_parameters()
    params.timeout = 3600
    ret = ctx.filecopy(params, str(space_usage_url), str('file://' + dest))
    if ret == 0:
        data_file = open(dest)
        data = json.load(data_file)
        data_file.close()
        if 'storageservice' in data:
            return consume_wlcg_json(data, agis_token)
        else:
            return consume_atlas_json(data, agis_token)


def get_space(rsename, protocol, client):
    retvalue = OK
    rse_settings = rsemgr.get_rse_info(rsename)
    if rse_settings['protocols']:
        schemes = [domain['scheme'] for domain in rse_settings['protocols']]
        if protocol == 'json':
            print 'Running on %s, using json method' % (rsename)
            try:
                totalsize, unusedsize = get_space_usage(rsename)
                freesize = long(unusedsize)
                totalsize = long(totalsize)
                usedsize = totalsize - freesize
            except Exception as error:
                print error
                retvalue = WARNING
        else:
            if 'srm' not in schemes:
                protocol = 'https'
            try:
                global_status, ret = rsemgr.get_space_usage(rse_settings, protocol)
                if global_status:
                    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 = WARNING
            except RSEProtocolNotSupported as error:
                print 'Protocol %s not found on %s' % (protocol, rsename)
                retvalue = WARNING
            except Exception as error:
                print "Failed to get rse(%s) space information: %s" % (rsename, str(error))
                retvalue = WARNING

        if retvalue == OK:
            print "Update RSE %s space usage (usedsize: %s, freesize: %s)" % (rsename, usedsize, freesize)
            client.set_rse_usage(rsename, protocol, usedsize, freesize)
            client.set_rse_usage(rsename, 'storage', usedsize, freesize)
            return retvalue
    return WARNING


if __name__ == "__main__":

    CLOUD = sys.argv[1]

    exitvalue = OK
    retvalue = OK
    usedsize = 0
    freesize = 0
    try:
        PROXY = config_get('nagios', 'proxy')
        os.environ["X509_USER_PROXY"] = PROXY
    except Exception as error:
        print "Failed to get proxy from rucio.cfg"

    CLIENT = Client()
    RSES = CLIENT.list_rses('cloud=%s' % CLOUD)
    for rse in sorted(RSES):
        rsename = rse['rse']
        print rsename
        rse_info = get_rse(rsename)
        if rse_info['availability_read']:
            attr = list_rse_attributes(rsename)
            if 'space_usage_method' in attr.keys():
                retvalue = get_space(rsename, protocol=attr['space_usage_method'], client=CLIENT)
            else:
                retvalue = get_space(rsename, protocol='srm', client=CLIENT)
        else:
            print '%s blacklisted for read. Skipping storage space collection' % rsename

        exitvalue = max(exitvalue, retvalue)
    sys.exit(exitvalue)
