#!/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-2019
#
# PY3K COMPATIBLE

from __future__ import print_function

import json
import sys
import traceback


import requests

from rucio.db.sqla.session import get_session
from rucio.common.exception import RucioException, RSENotFound
from rucio.core.rse import add_rse_attribute, get_rses_with_attribute, del_rse_attribute, get_rse_usage, get_rse
from rucio.core.rse_expression_parser import parse_expression


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)

    INTERNAL_PROXY_MAPPING = {}
    URL = 'http://atlas-agis-api.cern.ch/request/serviceobject/query/list/?json&state=ACTIVE'
    try:
        RESP = requests.get(url=URL)
        DIC = json.loads(RESP.content)
        for rse in DIC:
            for service in DIC[rse]:
                if 'protocols' in DIC[rse][service]:
                    for protocol in DIC[rse][service]['protocols']:
                        if DIC[rse][service]['protocols'][protocol]['doortype'] == 'proxyinternal':
                            INTERNAL_PROXY_MAPPING[rse] = DIC[rse][service]['protocols'][protocol]['endpoint']
    except Exception as error:
        print("Failed to load info from AGIS: %s" % str(error))
        sys.exit(WARNING)

    SPACE_USAGE = [rse['rse'] for rse in get_rses_with_attribute('space_usage_method')]
    HAS_PROXY = [rse['rse'] for rse in get_rses_with_attribute('root-proxy-internal')]
    NOT_FOR_EXTRA_COPY = [rse['rse'] for rse in get_rses_with_attribute('notforextracopy')]
    SCRATCHDISK_WITHOUT_SPACE = [rse['rse'] for rse in parse_expression(expression='type=SCRATCHDISK&freespace<50')]
    for rse in DATA:

        name = rse['name']

        session = get_session()
        # Check if RSE exists
        try:
            get_rse(name, session=session)
        except RSENotFound:
            continue

        print('Working on %s' % name)

        try:
            if rse['site'] in INTERNAL_PROXY_MAPPING:
                add_rse_attribute(name, 'root-proxy-internal', str(INTERNAL_PROXY_MAPPING[rse['site']]), session=session)
            elif name in HAS_PROXY:
                print('Will remove root-proxy-internal for site %s' % name)
                del_rse_attribute(name, 'root-proxy-internal', session=session)

            if name in SCRATCHDISK_WITHOUT_SPACE:
                add_rse_attribute(name, 'notforextracopy', 1, session=session)
            elif name in NOT_FOR_EXTRA_COPY:
                print('Will remove notforextracopy for site %s' % rse['name'])
                del_rse_attribute(name, 'notforextracopy', session=session)

            add_rse_attribute(name, 'ALL', '1', session=session)
            add_rse_attribute(name, 'tier', str(rse['tier_level']), session=session)
            add_rse_attribute(name, 'istape', str(rse['is_tape']), session=session)
            add_rse_attribute(name, 'cloud', str(rse['cloud']), session=session)
            add_rse_attribute(name, 'spacetoken', str(rse['token']), session=session)
            add_rse_attribute(name, 'site', str(rse['site']), session=session)
            add_rse_attribute(name, 'type', str(rse['type']), session=session)
            if str(rse['space_method']) == 'other':
                add_rse_attribute(name, 'space_usage_method', 'json', session=session)
            elif name in SPACE_USAGE:
                del_rse_attribute(name, 'space_usage_method', session=session)
            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, session=session)
            if rse['phys_groups']:
                add_rse_attribute(name, 'physgroup', str(rse['phys_groups'][0]), session=session)
            if rse['type'] not in ['OS_LOGS', 'OS_ES']:
                if isinstance(rse['servedrestfts']['MASTER'], list):
                    add_rse_attribute(name, 'fts', ','.join(rse['servedrestfts']['MASTER']), session=session)
                else:
                    add_rse_attribute(name, 'fts', str(rse['servedrestfts']['MASTER']), session=session)
            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']), session=session)
                else:
                    add_rse_attribute(name, 'fts_testing', str(rse['servedrestfts']['TESTING']), session=session)
            if 'datapolicies' in rse:
                add_rse_attribute(name, 'datapolicyt0disk', 'T0Disk' in rse['datapolicies'], session=session)
                add_rse_attribute(name, 'datapolicyt0tape', 'T0Tape' in rse['datapolicies'], session=session)
                add_rse_attribute(name, 'datapolicyt0taskoutput', 'T0TaskOutput' in rse['datapolicies'], session=session)
                add_rse_attribute(name, 'datapolicynucleus', 'Nucleus' in rse['datapolicies'], session=session)

            space_used = get_rse_usage(rse=name, source='storage', session=session)
            unavailable_space = get_rse_usage(rse=name, source='unavailable', session=session)
            expired = get_rse_usage(rse=name, source='expired', session=session)
            min_free_space = get_rse_usage(rse=name, source='min_free_space', session=session)
            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:
                min_free_space = min_free_space[0]['total'] if min_free_space else 0
                if space_used[0]['used'] == -1:
                    total_space = space_used[0]['total']
                    space_used = get_rse_usage(rse=name, source='rucio', session=session)
                    freespace = total_space - space_used[0]['used']
                else:
                    freespace = space_used[0]['free']
                freespace = float(freespace - unavailable_space + expired - min_free_space) / 1000 / 1000 / 1000 / 1000
                freespace = freespace if freespace > 0 else 0
                add_rse_attribute(name, 'freespace', int(freespace), session=session)

        except RucioException as error:
            print(error)
            sys.exit(CRITICAL)
        except Exception:
            print(traceback.format_exc())
            RESULT = WARNING
        finally:
            session.commit()
    sys.exit(RESULT)
