#!/usr/bin/env python
# Copyright 2012-2019 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:
# - Vincent Garonne <vincent.garonne@cern.ch>, 2013-2014
# - David Cameron <david.cameron@cern.ch>, 2014-2015,2020
# - Tomas Kouba <tomas.kouba@cern.ch>, 2014
# - Cedric Serfon <cedric.serfon@cern.ch>, 2016-2019
# - Dimitrios Christidis <dimitrios.christidis@cern.ch>, 2019
# - Mario Lassnig <mario.lassnig@cern.ch>, 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.common.config import config_get
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

# Takes DDM endpoint information from CRIC and adds selected attributes to RSEs
if __name__ == '__main__':
    try:
        PROXY = config_get('nagios', 'proxy')
    except Exception as error:
        print(error)
        print("Failed to get proxy from rucio.cfg")
        sys.exit(CRITICAL)

    URL = 'https://atlas-cric.cern.ch/api/atlas/ddmendpoint/query/?json'
    try:
        RESP = requests.get(url=URL, verify='/etc/pki/tls/certs/ca-bundle.crt', cert=PROXY)
        DATA = json.loads(RESP.content)
    except Exception as error:
        print("Failed to load info from CRIC: %s" % str(error))
        sys.exit(WARNING)

    INTERNAL_PROXY_MAPPING = {}
    URL = 'https://atlas-cric.cern.ch/api/core/service/query/?json'
    try:
        RESP = requests.get(url=URL, verify='/etc/pki/tls/certs/ca-bundle.crt', cert=PROXY)
        DIC = json.loads(RESP.content)
        for service in DIC:
            if 'protocols' in DIC[service]:
                for protocol in DIC[service]['protocols']:
                    if DIC[service]['protocols'][protocol]['doortype'] == 'proxyinternal':
                        INTERNAL_PROXY_MAPPING[DIC[service]['rcsite']] = DIC[service]['protocols'][protocol]['endpoint']
    except Exception as error:
        print("Failed to load info from CRIC: %s" % str(error))
        sys.exit(WARNING)

    SPACE_USAGE = [rse['rse'] for rse in get_rses_with_attribute('space_usage_method')]
    SOURCE_USED_SPACE = [rse['rse'] for rse in get_rses_with_attribute('source_for_used_space')]
    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 sorted(DATA.values(), key=lambda x: x['name']):

        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', True, 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, '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 rse['space_usage_url']:
                add_rse_attribute(name, 'srr_url', rse['space_usage_url'], session=session)
            elif name in SPACE_USAGE:
                add_rse_attribute(name, 'srr_url', rse['space_usage_url'], session=session)
            if rse['space_method'] == 'rucio':
                add_rse_attribute(name, 'source_for_used_space', 'rucio', session=session)
            elif name in SOURCE_USED_SPACE:
                del_rse_attribute(name, 'source_for_used_space', session=session)
            if rse['type'] == 'LOCALGROUPDISK' or rse['type'] == 'LOCALGROUPTAPE':
                country = rse['country_code'].lower()
                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 'MASTER' not in rse['servedrestfts']:
                    continue
                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 '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)
                add_rse_attribute(name, 'datapolicyanalysis', 'AnalysisData' in rse['datapolicies'], session=session)
            if rse['type'] in ('CALIBDISK', 'GROUPDISK', 'LOCALGROUPDISK'):
                add_rse_attribute(name, 'greedyDeletion', True, 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)
