#!/usr/bin/env python3
# Copyright 2012-2020 CERN for the benefit of the ATLAS collaboration.
#
# 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
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# 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
# - Mario Lassnig <mario.lassnig@cern.ch>, 2020

import json
import os
import sys

import gfal2

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_rses, list_rse_attributes, get_rse, set_rse_usage

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'

    """
    rse_attributes = list_rse_attributes(rse_name)
    space_usage_url = rse_attributes['srr_url']
    space_usage_key = rse_attributes['spacetoken']

    if space_usage_url not in CACHE:
        print(f'Downloading {space_usage_url}')
        ctx = gfal2.creat_context()
        ctx.set_opt_integer('CORE', 'NAMESPACE_TIMEOUT', 60)

        handle = ctx.open(space_usage_url, 'r')
        content = b''
        while True:
            chunk = handle.read_bytes(65000)
            if not chunk:
                break
            content += chunk

        CACHE[space_usage_url] = json.loads(content.decode())

    data = CACHE[space_usage_url]
    if 'storageservice' in data:
        return consume_wlcg_json(data, space_usage_key)
    else:
        return consume_atlas_json(data, space_usage_key)


def get_space(rsename, protocol):
    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(f'Running on {rsename}, using json method')
            try:
                totalsize, unusedsize = get_space_usage(rsename)
                freesize = int(unusedsize)
                totalsize = int(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 = int(ret["totalsize"])
                    freesize = int(ret["unusedsize"])
                    usedsize = totalsize - freesize
                else:
                    print(f'Failed to get rse({rsename}) space information: {ret!s}')
                    retvalue = WARNING
            except RSEProtocolNotSupported as error:
                print(f'Protocol {protocol} not found on {rsename}')
                retvalue = WARNING
            except Exception as error:
                print('Failed to get rse({rsename}) space information: {error!s}')
                retvalue = WARNING

        if retvalue == OK:
            print(f'Update RSE {rsename} space usage (usedsize: {usedsize}, freesize: {freesize})')
            set_rse_usage(rsename, 'storage', usedsize, freesize, 'root')
            return retvalue
    return WARNING


if __name__ == "__main__":

    CLOUD = sys.argv[1]
    CACHE = {}

    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')

    for rsename in sorted(rse['rse'] for rse in list_rses({'cloud': CLOUD})):
        print(rsename)
        rse_info = get_rse(rsename)
        if rse_info['availability_read']:
            if 'srr_url' in list_rse_attributes(rsename):
                retvalue = get_space(rsename, protocol='json')
            else:
                retvalue = get_space(rsename, protocol='srm')
        else:
            print(f'{rsename} excluded for read. Skipping storage space collection')

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