#!/usr/bin/env python
"""
 Copyright European Organization for Nuclear Research (CERN)

 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:
 - Vincent Garonne, <vincent.garonne@cern.ch>, 2013-2014
 - David Cameron, <david.cameron@cern.ch>, 2014-2015
 - Tomas Kouba, <tomas.kouba@cern.ch>, 2014
 - Cedric Serfon, <cedric.serfon@cern.ch>, 2016-2017
"""

import json
import sys
import traceback

import requests

from rucio.common.exception import RucioException, RSENotFound
from rucio.api.rse import add_rse_attribute, get_rses_with_attribute, del_rse_attribute, get_rse_usage, get_rse

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

RESULT = OK

# Map of countries to 2-letter code
# Eventually should be available in AGIS
COUNTRYCODES = {'Argentina': 'ar', 'Armenia': 'am', 'Australia': 'au', 'Austria': 'at',
                'Canada': 'ca', 'Switzerland': 'ch', 'Chile': 'cl', 'China': 'cn',
                'Czech Republic': 'cz', 'Germany': 'de', 'Denmark': 'dk', 'Spain': 'es',
                'France': 'fr', 'Greece': 'gr', 'Israel': 'il', 'Italy': 'it',
                'Japan': 'jp', 'Netherlands': 'nl', 'Nordic': 'dk', 'Norway': 'no',
                'Poland': 'pl', 'Portugal': 'pt', 'Romania': 'ro', 'Russian Federation': 'ru',
                'Sweden': 'se', 'Slovakia': 'sk', 'Slovenia': 'si', 'Turkey': 'tr', 'Taiwan': 'tw',
                'UK': 'uk', 'USA': 'us', 'South Africa': 'za'}


# Takes DDM endpoint information from AGIS and adds selected attributes to RSEs
if __name__ == '__main__':

    URL = 'http://atlas-agis-api.cern.ch/request/ddmendpoint/query/list/?json&state=ACTIVE&site_state=ACTIVE'
    try:
        RESP = requests.get(url=URL)
        DATA = json.loads(RESP.content)
    except Exception as error:
        print "Failed to load info from AGIS: %s" % str(error)
        sys.exit(WARNING)

    SPACE_USAGE = get_rses_with_attribute('space_usage_method')
    SPACE_USAGE = [rse['rse'] for rse in SPACE_USAGE]
    for rse in DATA:

        name = rse['name']

        # Check if RSE exists
        try:
            get_rse(name)
        except RSENotFound:
            continue

        print name

        try:
            add_rse_attribute(name, 'ALL', '1', 'root')
            add_rse_attribute(name, 'tier', str(rse['tier_level']), 'root')
            add_rse_attribute(name, 'istape', str(rse['is_tape']), 'root')
            add_rse_attribute(name, 'cloud', str(rse['cloud']), 'root')
            add_rse_attribute(name, 'spacetoken', str(rse['token']), 'root')
            add_rse_attribute(name, 'site', str(rse['site']), 'root')
            add_rse_attribute(name, 'type', str(rse['type']), 'root')
            if str(rse['space_method']) == 'other':
                add_rse_attribute(name, 'space_usage_method', 'json', 'root')
            elif name in SPACE_USAGE:
                del_rse_attribute(name, 'space_usage_method', 'root')
            if rse['type'] == 'LOCALGROUPDISK' or rse['type'] == 'LOCALGROUPTAPE':
                country = COUNTRYCODES[str(rse['country'])]
                if name.startswith('CERN'):
                    country = 'cern'
                add_rse_attribute(name, 'country', country, 'root')
            if rse['phys_groups']:
                add_rse_attribute(name, 'physgroup', str(rse['phys_groups'][0]), 'root')
            if rse['type'] not in ['OS_LOGS', 'OS_ES']:
                if isinstance(rse['servedrestfts']['MASTER'], list):
                    add_rse_attribute(name, 'fts', ','.join(rse['servedrestfts']['MASTER']), 'root')
                else:
                    add_rse_attribute(name, 'fts', str(rse['servedrestfts']['MASTER']), 'root')
            if rse['type'] not in ['OS_LOGS', 'OS_ES']:
                if isinstance(rse['servedrestfts']['TESTING'], list):
                    add_rse_attribute(name, 'fts_testing', ','.join(rse['servedrestfts']['TESTING']), 'root')
                else:
                    add_rse_attribute(name, 'fts_testing', str(rse['servedrestfts']['TESTING']), 'root')
            if 'datapolicies' in rse:
                add_rse_attribute(name, 'datapolicyt0disk', 'T0Disk' in rse['datapolicies'], 'root')
                add_rse_attribute(name, 'datapolicyt0tape', 'T0Tape' in rse['datapolicies'], 'root')
                add_rse_attribute(name, 'datapolicyt0taskoutput', 'T0TaskOutput' in rse['datapolicies'], 'root')
                add_rse_attribute(name, 'datapolicynucleus', 'Nucleus' in rse['datapolicies'], 'root')

            space_used = get_rse_usage(rse=name, issuer='root', source='storage')
            unavailable_space = get_rse_usage(rse=name, issuer='root', source='unavailable')
            expired = get_rse_usage(rse=name, issuer='root', source='expired')
            expired = expired[0]['total'] if expired != [] else 0
            if unavailable_space and unavailable_space[0]['total']:
                unavailable_space = unavailable_space[0]['total']
            else:
                unavailable_space = 0
            if space_used:
                if space_used[0]['used'] == -1:
                    total_space = space_used[0]['total']
                    space_used = get_rse_usage(rse=name, issuer='root', source='rucio')
                    freespace = total_space - space_used[0]['used']
                else:
                    freespace = space_used[0]['free']
                freespace = float(freespace - unavailable_space + expired) / 1000 / 1000 / 1000 / 1000
                freespace = freespace if freespace > 0 else 0
                add_rse_attribute(name, 'freespace', int(freespace), 'root')

        except RucioException as error:
            print str(error)
            sys.exit(CRITICAL)
        except:
            print traceback.format_exc()
            RESULT = WARNING
    sys.exit(RESULT)
