#!/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
# - Tomas Kouba, <tomas.kouba@cern.ch>, 2014

import json
import requests
import sys
import traceback

from rucio.common.exception import RucioException, RSENotFound
from rucio.api import rse as r

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

result = OK

# 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'
    try:
        resp = requests.get(url=url)
        data = json.loads(resp.content)
    except Exception, e:
        print "Failed to load info from AGIS: %s" % str(e)
        sys.exit(WARNING)

    for rse in data:

        # Only use active endpoints in AGIS
        if rse['state'] != 'ACTIVE':
            continue

        # Check if RSE exists
        try:
            r.get_rse(rse['name'])
        except RSENotFound:
            continue

        print rse['name']

        try:
            r.add_rse_attribute(rse['name'], 'ALL', '1', 'root')
            r.add_rse_attribute(rse['name'], 'tier', str(rse['tier_level']), 'root')
            r.add_rse_attribute(rse['name'], 'istape', str(rse['is_tape']), 'root')
            r.add_rse_attribute(rse['name'], 'cloud', str(rse['cloud']), 'root')
            r.add_rse_attribute(rse['name'], 'spacetoken', str(rse['token']), 'root')
            r.add_rse_attribute(rse['name'], 'country', str(rse['country']), 'root')
            r.add_rse_attribute(rse['name'], 'site', str(rse['site']), 'root')
            if rse['phys_groups']:
                r.add_rse_attribute(rse['name'], 'physgroup', str(rse['phys_groups'][0]), 'root')
            if isinstance(rse['servedrestfts']['MASTER'], list):
                r.add_rse_attribute(rse['name'], 'fts', ','.join(rse['servedrestfts']['MASTER']), 'root')
            else:
                r.add_rse_attribute(rse['name'], 'fts', str(rse['servedrestfts']['MASTER']), 'root')
            if isinstance(rse['servedrestfts']['TESTING'], list):
                r.add_rse_attribute(rse['name'], 'fts_testing', ','.join(rse['servedrestfts']['TESTING']), 'root')
            else:
                r.add_rse_attribute(rse['name'], 'fts_testing', str(rse['servedrestfts']['TESTING']), 'root')

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

        # TODO: set weights for data distribution here or in other collector

    sys.exit(result)
