#!/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-2017
 - David Cameron, <david.cameron@cern.ch>, 2015
"""

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

import requests

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 as 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 as 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 ['read_lan', 'read_wan', 'write_lan', 'write_wan', 'delete_lan', 'delete_wan']:
                    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)] = {'read_lan': 0, 'read_wan': 0, 'write_lan': 0, 'write_wan': 0, 'delete_lan': 0, 'delete_wan': 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]['read_lan'],
                                              "write": priority[prio]['write_lan'],
                                              "delete": priority[prio]['delete_lan']},
                                      "wan": {"read": priority[prio]['read_wan'],
                                              "write": priority[prio]['write_wan'],
                                              "delete": priority[prio]['delete_wan']}}}
                print params
                if impl:
                    try:
                        add_protocol(rse=rse['name'], parameter=params)
                    except Duplicate as 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 ['read_lan', 'read_wan', 'write_lan', 'write_wan', 'delete_lan', 'delete_wan']:
                    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)] = {'read_lan': 0, 'read_wan': 0, 'write_lan': 0, 'write_wan': 0, 'delete_lan': 0, 'delete_wan': 0}
                        priority[(o.scheme, o.netloc)][activity] = index
            if sum([act['read_lan'] for act in priority.values()]) == 0:
                for key in priority:
                    priority[key]['read_lan'] = priority[key]['read_wan']
            if sum([act['write_lan'] for act in priority.values()]) == 0:
                for key in priority:
                    priority[key]['write_lan'] = priority[key]['write_wan']
            if sum([act['delete_lan'] for act in priority.values()]) == 0:
                for key in priority:
                    priority[key]['delete_lan'] = priority[key]['delete_wan']

            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)]['read_lan'],
                                                  "write": priority[(o.scheme, o.netloc)]['write_lan'],
                                                  "delete": priority[(o.scheme, o.netloc)]['delete_lan']},
                                          "wan": {"read": priority[(o.scheme, o.netloc)]['read_wan'],
                                                  "write": priority[(o.scheme, o.netloc)]['write_wan'],
                                                  "delete": priority[(o.scheme, o.netloc)]['delete_wan']}}}

                    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:
                                for domain in ['lan', 'wan']:
                                    for act in ['read', 'write', 'delete']:
                                        if rucio_protocol['domains'][domain][act] != priority[(o.scheme, o.netloc)]['%s_%s' % (act, domain)]:
                                            print '%s : Protocol %s Activity %s_%s : priority in Rucio %s != priority in AGIS %s' % (rse['name'], o.scheme, act, domain,
                                                                                                                                     rucio_protocol['domains'][domain][act],
                                                                                                                                     priority[(o.scheme, o.netloc)]['%s_%s' % (act, domain)])
                                            update_protocols(rse['name'], o.scheme, {'domains': {domain: {act: priority[(o.scheme, o.netloc)]['%s_%s' % (act, domain)]}}}, hostname=netloc, port=port)
                            except RSEProtocolNotSupported as error:
                                print error
                        else:
                            print 'Will create protocol %s at %s with priorities read_lan,write_lan,delete_lan,read_wan,write_wan,delete_wan : ' \
                                  '%s,%s,%s,%s,%s,%s' % (o.scheme, rse['name'],
                                                         params['domains']['lan']['read'], params['domains']['lan']['write'], params['domains']['lan']['delete'],
                                                         params['domains']['wan']['read'], params['domains']['wan']['write'], params['domains']['wan']['delete'])
                            try:
                                add_protocol(rse=rse['name'], parameter=params)
                            except Exception as error:
                                print error
                except Duplicate as error:
                    pass
                except RSEProtocolPriorityError as 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)
