#!/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
# - Mario Lassnig, <mario.lassnig@cern.ch>, 2014
# - Cedric Serfon, <cedric.serfon@cern.ch>, 2014-2016
# - David Cameron, <david.cameron@cern.ch>, 2015

import json
import os.path
import requests
import sys
import traceback
import urlparse


from rucio.core.account import list_accounts, list_account_attributes
from rucio.core.account_limit import set_account_limit
from rucio.core.rse import get_rse_protocols, add_protocol, add_rse, update_protocols
from rucio.common.exception import Duplicate, RSEProtocolPriorityError, RSEProtocolNotSupported, RSENotFound


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

if __name__ == '__main__':

    URL = 'http://atlas-agis-api.cern.ch/request/ddmendpoint/query/list/?json&state=ACTIVE&site_state=ACTIVE'
    RESP = requests.get(url=URL)
    DATA = json.loads(RESP.content)
    RETVALUE = OK

    ACC_ATTRS = [(a['account'], list_account_attributes(a['account'])) for a in list_accounts()]
    ADMIN_ACCOUNTS = [x[0] for x in ACC_ATTRS if {'key': 'admin', 'value': True} in x[1]]
    for rse in DATA:
        print rse['name']
        try:
            deterministic = rse['is_deterministic']
            volatile = False
            add_rse(rse=rse['name'], deterministic=deterministic, volatile=volatile)
            # Set infinite quotas for admin accounts
            for account in ADMIN_ACCOUNTS:
                set_account_limit(account, rse['name'], -1)
            if not rse['is_tape']:
                set_account_limit('sonar', rse['name'], 10000000000000)
        except Duplicate, error:
            pass
        except Exception:
            RETVALUE = CRITICAL
            errno, errstr = sys.exc_info()[:2]
            trcbck = traceback.format_exc()
            print 'Interrupted processing with %s %s %s.' % (errno, errstr, trcbck)

        prefix = rse['endpoint']
        space_token = rse['token']

        try:
            rucio_protocols = get_rse_protocols(rse['name'], None)
        except RSENotFound, error:
            print error
            continue
        prot_read = []
        prot_write = []
        prot_delete = []
        priority = {}
        for activity in rse['aprotocols']:
            index = 0
            if activity in ['r', 'd', 'w']:
                for protocol, agis_prio, agis_prefix in rse['aprotocols'][activity]:
                    index += 1
                    o = urlparse.urlparse(protocol)
                    if o.scheme not in ('https', 'http', 'srm', 'gsiftp', 'root'):
                        continue
                    if (o.scheme, o.netloc) not in priority:
                        priority[(o.scheme, o.netloc)] = {'r': 0, 'w': 0, 'd': 0}
                    priority[(o.scheme, o.netloc)][activity] = index

        for protocol in rse['protocols']:
            try:
                o = urlparse.urlparse(protocol)
                if o.scheme not in ('https', 'http', 'srm', 'gsiftp', 'root'):
                    continue

                protocols = rse['protocols'][protocol]

                extended_attributes = None
                if o.scheme == 'srm':
                    extended_attributes = {"web_service_path": o.path + '?SFN=', "space_token": space_token}
                    impl = 'rucio.rse.protocols.gfal.Default'
                elif o.scheme == 'https' or o.scheme == 'http':
                    extended_attributes = None
                    impl = 'rucio.rse.protocols.webdav.Default'
                elif o.scheme == 'gsiftp':
                    extended_attributes = None
                    impl = 'rucio.rse.protocols.gsiftp.Default'
                elif o.scheme == 'root':
                    extended_attributes = None
                    impl = 'rucio.rse.protocols.xrootd.Default'
                else:
                    continue

                port = 8443
                netloc = o.netloc
                if o.port and str(o.port) in o.netloc:
                    netloc = o.netloc[:-len(':' + str(o.port))]
                    port = o.port
                else:
                    if o.scheme == 'https':
                        port = 443
                    elif o.scheme == 'gsiftp':
                        port = 2811
                    elif o.scheme == 'root':
                        port = 1094

                # For disk end-points not for tape
                prefix = rse['protocols'][protocol][0][2]
                if not rse['is_tape'] and not prefix.endswith('/rucio') and not prefix.endswith('/rucio/'):
                    prefix = os.path.join(prefix, 'rucio/')

                params = {'hostname': netloc,
                          'scheme': o.scheme,
                          'port': port,
                          'prefix': prefix,
                          'impl': impl,
                          'extended_attributes': extended_attributes,
                          'domains': {"lan": {"read": priority[(o.scheme, o.netloc)]['r'],
                                              "write": priority[(o.scheme, o.netloc)]['w'],
                                              "delete": priority[(o.scheme, o.netloc)]['d']},
                                      "wan": {"read": priority[(o.scheme, o.netloc)]['r'],
                                              "write": priority[(o.scheme, o.netloc)]['w'],
                                              "delete": priority[(o.scheme, o.netloc)]['d']}}}

                rucio_protocol = None
                for prot in rucio_protocols['protocols']:
                    if prot['scheme'] == o.scheme and prot['hostname'] == netloc and prot['port'] == port:
                        rucio_protocol = prot
                        break
                if params != rucio_protocol:
                    if rucio_protocol:
                        try:
                            activity_mapping = {'read': 'r', 'write': 'w', 'delete': 'd'}
                            for act in activity_mapping:
                                if rucio_protocol['domains']['lan'][act] != priority[(o.scheme, o.netloc)][activity_mapping[act]]:
                                    print '%s : Activity %s : priority in AGIS %s != priority in Rucio %s' % (rse['name'], act, rucio_protocol['domains']['lan'][act], priority[(o.scheme, o.netloc)][activity_mapping[act]])
                                    update_protocols(rse['name'], o.scheme, {'domains': {'lan': {act: priority[(o.scheme, o.netloc)][activity_mapping[act]]}}}, hostname=netloc, port=port)
                                    update_protocols(rse['name'], o.scheme, {'domains': {'wan': {act: priority[(o.scheme, o.netloc)][activity_mapping[act]]}}}, hostname=netloc, port=port)
                        except RSEProtocolNotSupported, error:
                            print error
                    else:
                        print 'Will create protocol %s at %s with priotities r,w,d : %s,%s,%s' % (o.scheme, rse['name'], params['domains']['lan']['read'], params['domains']['lan']['write'], params['domains']['lan']['delete'])
                        try:
                            add_protocol(rse=rse['name'], parameter=params)
                        except Exception, error:
                            print error
            except Duplicate, error:
                pass
            except RSEProtocolPriorityError, error:
                print 'RSE %s protocol %s: %s' % (rse['name'], o.scheme, error)
                if RETVALUE != CRITICAL:
                    RETVALUE = WARNING
            except Exception:
                RETVALUE = CRITICAL
                errno, errstr = sys.exc_info()[:2]
                trcbck = traceback.format_exc()
                print 'RSE %s protocol %s : Interrupted processing with %s %s %s.' % (rse['name'], o.scheme, errno, errstr, trcbck)
    sys.exit(RETVALUE)
