#!/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']
        deterministic = False
        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']

        existing_protocols = []
        try:
            rucio_protocols = get_rse_protocols(rse['name'], None)
            for prot in rucio_protocols['protocols']:
                existing_protocols.append((prot['scheme'], prot['hostname'], prot['port']))

        except RSENotFound, error:
            print error
            continue
        if rse['type'] in ['OS_ES', 'OS_LOGS']:
            print 'This is a Object store endpoint. Skipping the protocols'
            priority = {}
            for activity in rse['arprotocols']:
                index = 0
                if activity in ['r', 'd', 'w']:
                    for protocol in rse['arprotocols'][activity]:
                        index += 1
                        path = protocol['path']
                        o = urlparse.urlparse(protocol['endpoint'])
                        if (o.scheme, o.netloc, path) not in priority:
                            priority[(o.scheme, o.netloc, path)] = {'r': 0, 'w': 0, 'd': 0}
                        priority[(o.scheme, o.netloc, path)][activity] = index
            for prio in priority:
                scheme, host_with_port, prefix = prio
                if not prefix.endswith('/'):
                    prefix += '/'
                port = 443
                hostname = host_with_port
                if host_with_port.find(':') > -1:
                    hostname, port = host_with_port.split(':')
                impl = None
                if scheme == 's3':
                    impl = 'rucio.rse.protocols.s3boto.Default'
                elif scheme == 's3+rucio':
                    if rse['type'] == 'OS_ES':
                        impl = 'rucio.rse.protocols.ses3.Default'
                    else:
                        impl = 'rucio.rse.protocols.signeds3.Default'
                params = {'hostname': hostname,
                          'scheme': scheme,
                          'port': port,
                          'prefix': prefix,
                          'impl': impl,
                          'domains': {"lan": {"read": priority[prio]['r'],
                                              "write": priority[prio]['w'],
                                              "delete": priority[prio]['d']},
                                      "wan": {"read": priority[prio]['r'],
                                              "write": priority[prio]['w'],
                                              "delete": priority[prio]['d']}}}
                print params
                if impl:
                    try:
                        add_protocol(rse=rse['name'], parameter=params)
                    except Duplicate, error:
                        print error
                else:
                    print 'No implementation defined for %s on RSE %s' % (scheme, rse['name'])
                    RETVALUE = CRITICAL

        else:
            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', 'davs', 'dav'):
                            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', 'davs', 'dav'):
                        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 in ('davs', 'dav'):
                        extended_attributes = None
                        if rse['is_mkdir'] is True:
                            impl = 'rucio.rse.protocols.gfalv2.Default'
                        else:
                            impl = 'rucio.rse.protocols.gfal.Default'
                    elif o.scheme in ('https', 'http'):
                        extended_attributes = None
                        impl = 'rucio.rse.protocols.gfalv2.Default'
                    elif o.scheme == 'gsiftp':
                        extended_attributes = None
                        impl = 'rucio.rse.protocols.gfal.Default'
                    elif o.scheme == 'root':
                        extended_attributes = None
                        impl = 'rucio.rse.protocols.gfal.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 in ('https', 'davs'):
                            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 deterministic 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
                            try:
                                existing_protocols.remove((o.scheme, netloc, port))
                            except ValueError:
                                pass
                            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)
            if existing_protocols:
                RETVALUE = WARNING
                for scheme, hostname, port in existing_protocols:
                    print 'WARNING : Protocol %s://%s:%s is defined in Rucio but not in AGIS on RSE %s !!!!' % (scheme, hostname, port, rse['name'])
    sys.exit(RETVALUE)
