#!/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:
# - Mario Lassnig, <mario.lassnig@cern.ch>, 2012-2014
# - Vincent Garonne, <vincent.garonne@cern.ch>, 2012-2013
# - Thomas Beermann, <thomas.beermann@cern.ch>, 2012
# - Yun-Pin Sun, <yun-pin.sun@cern.ch>, 2013
# - Cedric Serfon <cedric.serfon@cern.ch>, 2013-2014
# - Martin Barisits <martin.barisits@cern.ch>, 2013-2014

"""
    Rucio CLI.
"""

import argcomplete
import argparse
import logging
import os
import random
import sys
import time


from rucio import client
from rucio.client import Client
from rucio import version
from rucio.client.accountclient import AccountClient
from rucio.client.didclient import DIDClient
from rucio.client.replicaclient import ReplicaClient
from rucio.client.metaclient import MetaClient
from rucio.client.pingclient import PingClient
from rucio.client.rseclient import RSEClient
from rucio.client.ruleclient import RuleClient
from rucio.client.scopeclient import ScopeClient
from rucio.client.subscriptionclient import SubscriptionClient
from rucio.common.exception import DataIdentifierAlreadyExists, Duplicate, FileAlreadyExists, AccessDenied, RessourceTemporaryUnavailable, DataIdentifierNotFound
from rucio.common.utils import adler32, generate_uuid
from rucio.rse import rsemanager as rsemgr

SUCCESS = 0
FAILURE = 1

DEFAULT_SECURE_PORT = 443
DEFAULT_PORT = 80


logger = logging.getLogger("user")


def setup_logger(logger):
    logger.setLevel(logging.DEBUG)
    hdlr = logging.StreamHandler()

    def emit_decorator(fn):
        def func(*args):
            levelno = args[0].levelno
            if(levelno >= logging.CRITICAL):
                color = '\033[31;1m'
            elif(levelno >= logging.ERROR):
                color = '\033[31;1m'
            elif(levelno >= logging.WARNING):
                color = '\033[33;1m'
            elif(levelno >= logging.INFO):
                color = '\033[32;1m'
            elif(levelno >= logging.DEBUG):
                color = '\033[36;1m'
            else:
                color = '\033[0m'
            formatter = logging.Formatter('{0}%(asctime)s %(levelname)s [%(message)s]\033[0m'.format(color))
            hdlr.setFormatter(formatter)
            return fn(*args)
        return func
    hdlr.emit = emit_decorator(hdlr.emit)
    logger.addHandler(hdlr)

setup_logger(logger)


def get_client(args):
    """
    Returns a new client object.
    specified by the --host and --port options
    supplied to the CLI
    """

    return client.Client(host="0.0.0.0")


def ping(args):
    """
    Pings a Rucio server.
    """

    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}
    else:
        creds = None

    client = PingClient(rucio_host=args.host, auth_host=args.auth_host,
                        account=args.account,
                        auth_type=args.auth_strategy, creds=creds,
                        ca_cert=args.ca_certificate, timeout=args.timeout)

    server_info = client.ping()
    if server_info:
        print server_info['version']
        return SUCCESS
    print 'Ping failed'
    return FAILURE


def whoami_account(args):
    """
    %(prog)s show [options] <field1=value1 field2=value2 ...>

    Show extended information of a given account
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    accountClient = AccountClient(rucio_host=args.host, auth_host=args.auth_host,
                                  account=args.account,
                                  auth_type=args.auth_strategy, creds=creds,
                                  ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        info = accountClient.whoami()
    except Exception, e:
            print 'Failed to show account'
            print e
            return FAILURE

    for k in info:
        print k.ljust(10) + ' : ' + str(info[k])

    return SUCCESS


def add_file(args):
    """
    %(prog)s show [options] <field1=value1 field2=value2 ...>

    Add file.
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    scope = args.scope
    lfn = args.lfn
    if scope is None:
        if args.lfn.find(':') == -1:
            print 'Error: Cannot extract scope from %s (notation: scope:filename) !' % args.lfn
            return FAILURE
        scope, lfn = args.lfn.split(':')

    client = RSEClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account,
                       auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)
    try:
        client.add_file(rse=args.rse, scope=scope, lfn=lfn)
        print 'Added new file replica: %s-%s' % (args.rse, args.lfn)
    except Exception, e:
            print 'Failed to add file replica'
            print e
            return FAILURE
    return SUCCESS


def list_replicas(args):
    """
    %(prog)s list [options] <field1=value1 field2=value2 ...>

    List file replicas
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = ReplicaClient(rucio_host=args.host, auth_host=args.auth_host,
                           account=args.account, auth_type=args.auth_strategy, creds=creds,
                           ca_cert=args.ca_certificate, timeout=args.timeout)

    protocols = None
    if args.protocols:
        protocols = args.protocols.split(',')

    dids = []
    try:
        for did in args.dids:
            try:
                scope, name = did.split(':')
            except ValueError, e:
                print 'ERROR cannot extract the scope and name from %s : [%s]' % (did, e)
                return FAILURE
            dids.append({'scope': scope, 'name': name})
        replicas = client.list_replicas(dids, schemes=protocols)
        print 'Scope\tName\t\t\tFilesize\tadler32\tReplicas'
        for replica in replicas:
            if 'bytes' in replica:
                str = '%s\t%s\t%s\t%s\t' % (replica['scope'], replica['name'], replica['bytes'], replica['adler32'])
                for rse in replica['rses']:
                    line = '%s%s\t:\t' % (str, rse)
                    for pfn in replica['rses'][rse]:
                        output = '%s%s' % (line, pfn)
                        print output.expandtabs()
                        line = len(line.expandtabs()) * ' '
                        str = len(str.expandtabs()) * ' '
            else:
                str = '%s\t%s\t%s\t%s\t%s' % (replica['scope'], replica['name'], '???', '???', 'Unavalaible')
                print str.expandtabs()
    except Exception, e:
        print 'Failed to list file replicas for %(scope)s:%(name)s' % locals()
        print e


def add_dataset(args):
    """
    %(prog)s add-dataset [options] <dsn>

    Add a dataset identifier.
    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    try:
        scope, name = args.did.split(':')
    except ValueError, e:
        print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.did, e)
        return FAILURE

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)
    try:
        client.add_dataset(scope=scope, name=name, statuses={'monotonic': args.monotonic})
        print 'Added %s:%s' % (scope, name)
    except Exception, e:
        print 'Failed to add data identifier'
        print e
        return FAILURE
    return SUCCESS


def add_container(args):
    """
    %(prog)s add-container [options] <dsn>

    Add a container identifier.
    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    try:
        scope, name = args.did.split(':')
    except ValueError, e:
        print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.did, e)
        return FAILURE

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)
    try:
        client.add_container(scope=scope, name=name, statuses={'monotonic': args.monotonic})
        print 'Added %s:%s' % (scope, name)
    except Exception, e:
        print 'Failed to add data identifier'
        print e
        return FAILURE
    return SUCCESS


def attach(args):
    """
    %(prog)s attach [options] <field1=value1 field2=value2 ...>

    Attach a data identifier.
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        dids = [{'scope': did.split(':')[0], 'name': did.split(':')[1]} for did in args.dids]
        scope, name = args.todid.split(':')
    except ValueError, e:
        logger.error('ERROR cannot extract the scope and name from {0} : [{1}]'.format(args.todid, e))
        return FAILURE
    except Exception, e:
        logger.warning('Unknown exception occurred : {0}'.format(e))

    try:
        client.attach_dids(scope=scope, name=name, dids=dids)
    except Exception, e:
        logger.error('Failed to attach data identifier\n{0}'.format(e))
        return FAILURE
    return SUCCESS


def detach(args):
    """
    %(prog)s detach [options] <field1=value1 field2=value2 ...>

    Detach data identifier.
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        dids = [{'scope': did.split(':')[0], 'name': did.split(':')[1]} for did in args.dids]
        scope, name = args.fromdid.split(':')
    except ValueError, e:
        logger.error('ERROR cannot extract the scope and name from {0} : [{1}]'.format(args.fromdid, e))
        return FAILURE
    except Exception, e:
        logger.warning('Unknown exception occurred : {0}'.format(e))

    try:
        client.detach_dids(scope=scope, name=name, dids=dids)
    except Exception, e:
        logger.error('Failed to detach data identifier\n{0}'.format(e))
        return FAILURE
    return SUCCESS


def add_files_to_dataset(args):
    """
    %(prog)s add-files-to-dataset [options] <field1=value1 field2=value2 ...>

    Attach files to a dataset.
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        dids = [{'scope': did.split(':')[0], 'name': did.split(':')[1]} for did in args.files]
        scope, name = args.todid.split(':')
    except ValueError, e:
        print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.todid, e)
        return FAILURE
    except Exception, e:
        print e

    try:
        client.add_files_to_dataset(scope=scope, name=name, files=dids)
    except Exception, e:
        print 'Failed to attach files to the dataset'
        print e
        return FAILURE
    return SUCCESS


def add_datasets_to_container(args):
    """
    %(prog)s add-datasets-to-container [options] <field1=value1 field2=value2 ...>

    Attach datasets to a container.
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        dids = [{'scope': did.split(':')[0], 'name': did.split(':')[1]} for did in args.datasets]
        scope, name = args.todid.split(':')
    except ValueError, e:
        print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.todid, e)
        return FAILURE
    except Exception, e:
        print e

    try:
        client.add_datasets_to_container(scope=scope, name=name, dsns=dids)
    except Exception, e:
        print 'Failed to attach datasets to the container'
        print e
        return FAILURE
    return SUCCESS


def add_containers_to_container(args):
    """
    %(prog)s add-containers-to-container [options] <field1=value1 field2=value2 ...>

    Attach containers to a container.
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        dids = [{'scope': did.split(':')[0], 'name': did.split(':')[1]} for did in args.containers]
        scope, name = args.todid.split(':')
    except ValueError, e:
        print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.todid, e)
        return FAILURE
    except Exception, e:
        print e

    try:
        client.add_containers_to_container(scope=scope, name=name, cnts=dids)
    except Exception, e:
        print 'Failed to attach containers to the container'
        print e
        return FAILURE
    return SUCCESS


def list_dids(args):
    """
    %(prog)s list <did_1> ... <did_n>

    List data identifier contents.
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    if not args.dids:
        scope = args.scope
        if not args.scope:  # list scopes
            client = ScopeClient(rucio_host=args.host, auth_host=args.auth_host,
                                 account=args.account,
                                 auth_type=args.auth_strategy, creds=creds,
                                 ca_cert=args.ca_certificate, timeout=args.timeout)
            try:
                account = args.account
                if not account:
                    accountClient = AccountClient(rucio_host=args.host, auth_host=args.auth_host,
                                                  account=args.account,
                                                  auth_type=args.auth_strategy, creds=creds,
                                                  ca_cert=args.ca_certificate, timeout=args.timeout)
                    info = accountClient.whoami()
                    account = info['account']
                scopes = client.list_scopes_for_account(account)
                if len(scopes) > 1:
                    print 'Too many scopes to list. Please specify one scope'
                    return FAILURE
                scope = scopes[0]
            except Exception, e:
                print 'Failed to list scopes'
                print e
                return FAILURE
        client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                           account=args.account, auth_type=args.auth_strategy, creds=creds,
                           ca_cert=args.ca_certificate, timeout=args.timeout)
        try:
            dids = client.scope_list(scope=scope, recursive=args.recursive)
            for d in dids:
                if d['level'] == 0:
                    print '%(scope)s:%(name)s [%(type)s]' % d
                else:
                    print '|    ' * d['level'] + '|- %(scope)s:%(name)s [%(type)s]' % d
        except Exception, e:
            print e
            return FAILURE
    else:
        client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                           account=args.account, auth_type=args.auth_strategy, creds=creds,
                           ca_cert=args.ca_certificate, timeout=args.timeout)
        for did in args.dids:
            try:
                scope, name = did.split(':')
                dids = client.scope_list(scope=scope, name=name, recursive=args.recursive)
                for d in dids:
                    if d['level'] == 0:
                        print '%(scope)s:%(name)s [%(type)s]' % d
                    else:
                        print '|    ' * d['level'] + '|- %(scope)s:%(name)s [%(type)s]' % d
            except ValueError, e:
                print 'ERROR cannot extract the scope and name from %s : [%s]' % (did, e)
                return FAILURE
            except Exception, e:
                print e

    return SUCCESS


def list_scopes(args):
    """
    %(prog)s list-scopes <scope>

    List scopes.
    """
    # For the moment..
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = ScopeClient(rucio_host=args.host, auth_host=args.auth_host,
                         account=args.account,
                         auth_type=args.auth_strategy, creds=creds,
                         ca_cert=args.ca_certificate, timeout=args.timeout)
    try:
        scopes = client.list_scopes()
    except Exception, e:
            print 'Failed to list scopes'
            print e
            return FAILURE

    for scope in scopes:
        print scope

    return SUCCESS


def scope_list(args):
    """
    %(prog)s scope_list <scope>

    List all data identifiers in given scope.
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        dids = client.scope_list(scope=args.scope)
        for d in dids:
            print '%(scope)s:%(name)s [%(type)s]' % d
    except Exception, e:
        print e
        return FAILURE

    return SUCCESS


def list_files(args):
    """
    %(prog)s list-files [options] <field1=value1 field2=value2 ...>

    List data identifier contents.
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    for did in args.dids:
        try:
            scope, name = did.split(':')
            for file in client.list_files(scope=scope, name=name):
                print '%(scope)s:%(name)s' % file  # %(size)s\t%(checksum)s
        except ValueError, e:
            print 'ERROR cannot extract the scope and name from %s : [%s]' % (did, e)
            return FAILURE
        except Exception, e:
            print e

    return SUCCESS


def close(args):
    """
    %(prog)s close [options] <field1=value1 field2=value2 ...>

    Close data identifier.
    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)
    for did in args.dids:
        try:
            scope, name = did.split(':')
            client.set_status(scope=scope, name=name, open=False)
            print '%(scope)s:%(name)s has been closed.' % locals()
        except ValueError, e:
            print 'ERROR cannot extract the scope and name from %s : [%s]' % (did, e)
            return FAILURE
        except Exception, e:
            print e

    return SUCCESS


def stat(args):
    """
    %(prog)s stat [options] <field1=value1 field2=value2 ...>

    List attributes and statuses about data identifiers..
    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)
    for did in args.dids:
        try:
            scope, name = did.split(':')
            info = client.get_did(scope=scope, name=name)
            for k, v in info.iteritems():
                print '%(k)s: %(v)s' % locals()
        except ValueError, e:
            print 'ERROR cannot extract the scope and name from %s : [%s]' % (did, e)
            return FAILURE
        except Exception, e:
            print e

    return SUCCESS


def delete(args):
    """
    %(prog)s delete [options] <field1=value1 field2=value2 ...>

    Delete data identifier.
    """
    raise NotImplementedError('operation not implemented (delete).')


def upload(args):
    """
    %(prog)s upload [options] <field1=value1 field2=value2 ...>

    Upload files into Rucio
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = Client(rucio_host=args.host, auth_host=args.auth_host,
                    account=args.account, auth_type=args.auth_strategy, creds=creds,
                    ca_cert=args.ca_certificate, timeout=args.timeout)

    files = args.files
    list_files = []
    files_to_list = []
    lfns = {}
    revert_dict = {}
    logger.debug('Looping over the files')
    for name in files:
        try:
            size = os.stat(name).st_size
            checksum = adler32(name)
            logger.debug('Extracting filesize (%s) and checksum (%s) for file %s:%s' % (str(size), checksum, args.scope, os.path.basename(name)))
            files_to_list.append({'scope': args.scope, 'name': os.path.basename(name)})
            list_files.append({'scope': args.scope, 'name': os.path.basename(name), 'bytes': size, 'adler32': checksum, 'state': 'C', 'meta': {'guid': generate_uuid()}})
            if not os.path.dirname(name) in lfns:
                lfns[os.path.dirname(name)] = []
            lfns[os.path.dirname(name)].append({'name': os.path.basename(name), 'scope': args.scope, 'adler32': checksum, 'filesize': size})
            revert_dict[args.scope, os.path.basename(name)] = os.path.dirname(name)
        except OSError, e:
            logger.error(e)
            logger.error("No operation will be performed. Exiting!")
            return FAILURE

    metadata = {}
    if args.project:
        metadata['project'] = args.project
    if args.datatype:
        metadata['datatype'] = args.datatype
    if args.run_number:
        metadata['run_number'] = args.run_number
    if args.stream_name:
        metadata['stream_name'] = args.stream_name
    if args.prod_step:
        metadata['prod_step'] = args.prod_step
    if args.version:
        metadata['version'] = args.version
    if args.campaign:
        metadata['campaign'] = args.campaign

    rse_settings = rsemgr.get_rse_info(args.rse)
    if rse_settings['availability_write'] != 1:
        logger.critical('RSE is not available for write now')
        return FAILURE

    if args.account is None:
        account = client.whoami()['account']
    else:
        account = args.account
    logger.debug('Using account %s' % (account))

    if args.did:
        try:
            dsscope, dsname = args.did.split(':')
            meta = client.get_metadata(scope=dsscope, name=dsname)
            # Dataset already exist
        except ValueError, e:
            logger.error('ERROR cannot extract the scope and name from %s : [%s]' % (args.did, e))
            return FAILURE
        except DataIdentifierNotFound:
            # Dataset non existent. Must be created.
            client.add_dataset(scope=dsscope, name=dsname, rules=[{'account': account, 'copies': 1, 'rse_expression': args.rse, 'grouping': 'DATASET'}], meta=metadata)
            logger.info('Dataset successfully created')
    else:
        # no dataset provided, only upload files
        dsscope = None
        dsname = None

    # Adding files to the catalog
    for f in list_files:
        try:  # If the did already exist in the catalog, only shold be upload if the checksum is the same
            meta = client.get_metadata(f['scope'], f['name'])
            logger.warning("The file {0}:{1} already exist in the catalog and will not be added.".format(f['scope'], f['name']))
            if rsemgr.exists(rse_settings=rse_settings, files={'name': f['name'], 'scope': f['scope']}):
                logger.warning('File {0}:{1} already exists on RSE. Will not try to reupload'.format(f['scope'], f['name']))
            else:
                if meta['adler32'] == f['adler32']:
                    logger.info('Local files and file %s:%s recorded in Rucio have the same checksum. Will try the upload' % (f['scope'], f['name']))
                    directory = revert_dict[f['scope'], f['name']]
                    rsemgr.upload(rse_settings=rse_settings, lfns=[{'name': f['name'], 'scope': f['scope'], 'adler32': f['adler32'], 'filesize': f['bytes']}], source_dir=directory)
                    logger.info('File %s:%s successfully uploaded on the storage' % (f['scope'], f['name']))
                else:
                    raise DataIdentifierAlreadyExists

        except DataIdentifierNotFound:
            try:
                logger.info('Adding replicas in Rucio catalog')
                client.add_replicas(files=[f], rse=args.rse)
                logger.info('Replicas successfully added')
                logger.info('Adding replication rule on RSE {0} for the file {1}:{2}'.format(args.rse, f['scope'], f['name']))
                client.add_replication_rule([f], copies=1, rse_expression=args.rse)
                directory = revert_dict[f['scope'], f['name']]
                rsemgr.upload(rse_settings=rse_settings, lfns=[{'name': f['name'], 'scope': f['scope'], 'adler32': f['adler32'], 'filesize': f['bytes']}], source_dir=directory)
                logger.info('File {0}:{1} successfully uploaded on the storage'.format(f['scope'], f['name']))
            except (Duplicate, FileAlreadyExists), e:
                logger.warning(e)
                return FAILURE
            except RessourceTemporaryUnavailable, e:
                logger.error(e)
                return FAILURE
        except DataIdentifierAlreadyExists, e:
            logger.error(e)
        if dsname:
            # A dataset is provided. Must add the files to the dataset.
            for f in list_files:
                try:
                    client.add_files_to_dataset(scope=dsscope, name=dsname, files=[f])
                except Exception, e:
                    logger.warning('Failed to attach file {0} to the dataset'.format(f))
                    logger.warning(e)
                    logger.warning("Continuing with the next one")

    replicas = []
    replica_dictionary = {}
    for rep in client.list_replicas(files_to_list):
        replica_dictionary[rep['scope'], rep['name']] = rep['rses'].keys()
    for file in list_files:
        if (file['scope'], file['name']) not in replica_dictionary:
            file['state'] = 'A'
            replicas.append(file)
        elif args.rse not in replica_dictionary[file['scope'], file['name']]:
            file['state'] = 'A'
            replicas.append(file)
    if replicas != []:
        logger.info('Will update the file replicas states')
        try:
            client.update_replicas_states(rse=args.rse, files=replicas)
        except AccessDenied, e:
            logger.info(e)
            return FAILURE
        logger.info('File replicas states successfully updated')
    return SUCCESS


def download(args):
    """
    %(prog)s download [options] <field1=value1 field2=value2 ...>

    Download files from Rucio
    """

    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = ReplicaClient(rucio_host=args.host, auth_host=args.auth_host,
                           account=args.account, auth_type=args.auth_strategy, creds=creds,
                           ca_cert=args.ca_certificate, timeout=args.timeout)
    rse_dict = {}
    summary = {}
    for did in args.dids:
        try:
            scope, name = did.split(':')
            logger.info('Starting download for %s:%s' % (scope, name))
            summary['%s:%s' % (scope, name)] = {}
            logger.debug('Getting the list of replicas')
            files = client.list_replicas([{'scope': scope, 'name': name}])
            for file in files:
                if os.path.isfile('%s/%s/%s' % (args.dir, file['scope'], file['name'])):
                    logger.info('File %s:%s already exists locally' % (file['scope'], file['name']))
                    summary['%s:%s' % (scope, name)]['%s:%s' % (file['scope'], file['name'])] = 2
                else:
                    summary['%s:%s' % (scope, name)]['%s:%s' % (file['scope'], file['name'])] = 0
                    list_rses = file['rses'].keys()
                    random.shuffle(list_rses)
                    logger.debug('Choosing RSE')
                    index = 0
                    for rse in list_rses:
                        if rse not in rse_dict:
                            rse_dict[rse] = rsemgr.get_rse_info(rse)
                        if rse_dict[rse]['availability_read']:
                            logger.debug('Getting file %s:%s from %s' % (file['scope'], file['name'], rse))
                            try:
                                rsemgr.download(rse_dict[rse], files=[{'name': file['name'], 'scope': file['scope'], 'adler32': file['adler32']}, ], dest_dir=args.dir, printstatements=True)
                                logger.info('File %s:%s successfully downloaded from %s' % (file['scope'], file['name'], rse))
                                summary['%s:%s' % (scope, name)]['%s:%s' % (file['scope'], file['name'])] = 1
                                break
                            except Exception, e:
                                logger.warning(e)
                                index += 1
                                if index != len(rse):
                                    logger.debug('Will retry download on an other RSE')
                    if summary['%s:%s' % (scope, name)]['%s:%s' % (file['scope'], file['name'])] == 0:
                        logger.error('Cannot download file %s:%s' % (file['scope'], file['name']))
            logger.info('Download operation for %s:%s done' % (scope, name))
        except ValueError, e:
            logger.error('ERROR cannot extract the scope and name from %s : [%s]' % (did, e))
            return FAILURE
        except Exception, e:
            logger.error('Failed to download %(scope)s:%(name)s' % locals())
            logger.error(e)
    print '----------------------------------'
    print 'Download summary'
    for did in summary:
        print 'DID %s' % (did)
        total_files = len(summary[did])
        downloaded_files = 0
        local_files = 0
        for file in summary[did]:
            if summary[did][file] == 1:
                downloaded_files += 1
            elif summary[did][file] == 2:
                local_files += 1


def get_metadata(args):
    """
    %(prog)s get_metadata [options] <field1=value1 field2=value2 ...>

    Get data identifier metadata
    """
    # For the moment..
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    for did in args.dids:
        try:
            scope, name = did.split(':')
            meta = client.get_metadata(scope=scope, name=name)
            for k in meta:
                print '%s: %s' % (k, meta[k])
        except ValueError, e:
            print 'ERROR cannot extract the scope and name from %s : [%s]' % (did, e)
            return FAILURE
        except Exception, e:
            print 'Failed to get metadata for %(scope)s:%(name)s' % locals()
            print e


def set_metadata(args):
    """
    %(prog)s set_metadata [options] <field1=value1 field2=value2 ...>

    Set data identifier metadata
    """
    # For the moment..
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account, auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        client.set_metadata(scope=args.did.split(':')[0], name=args.did.split(':')[1], key=args.key, value=args.value)
    except ValueError, e:
        print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.did, e)
        return FAILURE
    except Exception, e:
        print 'Failed to set metadata'
        print e
        return FAILURE


def delete_metadata(args):
    """
    %(prog)s set_metadata [options] <field1=value1 field2=value2 ...>

    Delete data identifier metadata
    """
    # For the moment..
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    MetaClient(rucio_host=args.host, auth_host=args.auth_host,
               account=args.account,
               auth_type=args.auth_strategy, creds=creds,
               ca_cert=args.ca_certificate, timeout=args.timeout)


def add_rule(args):
    """
    %(prog)s add-rule <did> <copies> <rse-expression> [options]

    Add a rule to a did.
    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    dids = []
    for did in args.dids:
        try:
            scope, name = did.split(':')
        except ValueError, e:
            print 'ERROR cannot extract the scope and name from %s : [%s]' % (did, e)
            return FAILURE
        dids.append({'scope': scope, 'name': name})

    client = RuleClient(rucio_host=args.host, auth_host=args.auth_host,
                        account=args.account, auth_type=args.auth_strategy, creds=creds,
                        ca_cert=args.ca_certificate, timeout=args.timeout)
    try:
        rule_ids = client.add_replication_rule(dids=dids,
                                               copies=args.copies,
                                               rse_expression=args.rse_expression,
                                               weight=args.weight,
                                               lifetime=args.lifetime,
                                               grouping=args.grouping,
                                               account=args.account,
                                               locked=args.locked,
                                               source_replica_expression=args.source_replica_expression,
                                               notify=args.notify)
        for rule in rule_ids:
            print rule
    except Exception, e:
        print 'Failed to add rule to data identifier'
        print e
        return FAILURE
    return SUCCESS


def delete_rule(args):
    """
    %(prog)s delete-rule [options] <ruleid>

    Delete a rule.
    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = RuleClient(rucio_host=args.host, auth_host=args.auth_host,
                        account=args.account, auth_type=args.auth_strategy, creds=creds,
                        ca_cert=args.ca_certificate, timeout=args.timeout)
    try:
        client.delete_replication_rule(rule_id=args.rule_id)
        print 'Removed Rule'
    except Exception, e:
        print 'Failed to delete replication rule'
        print e
        return FAILURE
    return SUCCESS


def update_rule(args):
    """
    %(prog)s update-rule [options] <ruleid>

    Update a rule.
    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = RuleClient(rucio_host=args.host, auth_host=args.auth_host,
                        account=args.account, auth_type=args.auth_strategy, creds=creds,
                        ca_cert=args.ca_certificate, timeout=args.timeout)
    try:
        options = {}
        if args.lifetime:
            options['lifetime'] = None if args.lifetime == "None" else int(args.lifetime)
        if args.locked:
            if args.locked == "True":
                options['locked'] = True
            elif args.locked == "False":
                options['locked'] = False
        print options
        client.update_replication_rule(rule_id=args.rule_id, options=options)
        print 'Updated Rule'
    except Exception, e:
        print 'Failed to update replication rule'
        print e
        return FAILURE
    return SUCCESS


def info_rule(args):
    """
    %(prog)s rule-info [options] <ruleid>

    Retrieve information about a rule.
    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = RuleClient(rucio_host=args.host, auth_host=args.auth_host,
                        account=args.account, auth_type=args.auth_strategy, creds=creds,
                        ca_cert=args.ca_certificate, timeout=args.timeout)
    try:
        rule = client.get_replication_rule(rule_id=args.rule_id)
        print "Id:                         %s" % rule['id']
        print "Account:                    %s" % rule['account']
        print "Scope:                      %s" % rule['scope']
        print "Name:                       %s" % rule['name']
        print "RSE Expression:             %s" % rule['rse_expression']
        print "Copies:                     %s" % rule['copies']
        print "State:                      %s" % rule['state']
        print "Locks OK/REPLICATING/STUCK: %s/%s/%s" % (rule['locks_ok_cnt'], rule['locks_replicating_cnt'], rule['locks_stuck_cnt'])
        print "Grouping:                   %s" % rule['grouping']
        print "Expires at:                 %s" % rule['expires_at']
        print "Locked:                     %s" % rule['locked']
        print "Weight:                     %s" % rule['weight']
        print "Created at:                 %s" % rule['created_at']
        print "Updated at:                 %s" % rule['updated_at']
        print "Error:                      %s" % rule['error']
        print "Subscription Id:            %s" % rule['subscription_id']
        print "Source replica expression:  %s" % rule['source_replica_expression']
    except Exception, e:
        print 'Failed to get replication rule'
        print e
        return FAILURE
    return SUCCESS


def list_rules(args):
    """
    %(prog)s list-rules ...

    List rules.
    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    if args.rule_id:
        client = RuleClient(rucio_host=args.host, auth_host=args.auth_host,
                            account=args.account,
                            auth_type=args.auth_strategy, creds=creds,
                            ca_cert=args.ca_certificate, timeout=args.timeout)
        try:
            rules = [client.get_replication_rule(args.rule_id)]
        except Exception, e:
            logger.error('Failed to list rules.\n{0}'.format(e))
            return FAILURE
    elif args.file:
        client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                           account=args.account,
                           auth_type=args.auth_strategy, creds=creds,
                           ca_cert=args.ca_certificate, timeout=args.timeout)
        try:
            scope, name = args.file.split(':')
            rules = client.list_associated_rules_for_file(scope=scope, name=name)
        except ValueError, e:
            print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.did, e)
            return FAILURE
        except Exception, e:
            print 'Failed to list rules'
            print e
            return FAILURE
    elif args.did:
        client = DIDClient(rucio_host=args.host, auth_host=args.auth_host,
                           account=args.account,
                           auth_type=args.auth_strategy, creds=creds,
                           ca_cert=args.ca_certificate, timeout=args.timeout)
        try:
            scope, name = args.did.split(':')
            client.get_metadata(scope=scope, name=name)
            rules = client.list_did_rules(scope=scope, name=name)
        except ValueError, e:
            logger.error('ERROR cannot extract the scope and name from {0} : [{1}]'.format(args.did, e))
            return FAILURE
        except Exception, e:
            logger.error('Failed to list rules.\n{0}'.format(e))
            return FAILURE
    elif args.rule_account:
        client = AccountClient(rucio_host=args.host, auth_host=args.auth_host,
                               account=args.account,
                               auth_type=args.auth_strategy, creds=creds,
                               ca_cert=args.ca_certificate, timeout=args.timeout)
        try:
            rules = client.list_account_rules(account=args.rule_account)
        except Exception, e:
            print 'Failed to list rules'
            print e
            return FAILURE
    elif args.subscription:
        account = args.subscription[0]
        name = args.subscription[1]
        client = SubscriptionClient(rucio_host=args.host, auth_host=args.auth_host,
                                    account=args.account,
                                    auth_type=args.auth_strategy, creds=creds,
                                    ca_cert=args.ca_certificate, timeout=args.timeout)
        try:
            rules = client.list_subscription_rules(account=account, name=name)
        except Exception, e:
            logger.error('Failed to list rules\n{0}'.format(e))
            return FAILURE
    else:
        print 'At least one option has to be given. Use -h to list the options.'
        return FAILURE

    print "ID (account) SCOPE:NAME: STATE [LOCKS_OK/REPLICATING/STUCK], RSE_EXPRESSION, COPIES"
    print "==================================================================================="
    for rule in rules:
        print "%s (%s) %s:%s: %s[%s/%s/%s], \"%s\", %s " % (rule['id'], rule['account'], rule['scope'], rule['name'], rule['state'], rule['locks_ok_cnt'], rule['locks_replicating_cnt'], rule['locks_stuck_cnt'], rule['rse_expression'], rule['copies'])

    return SUCCESS


def list_rses(args):
    """
    %(prog)s list [options] <field1=value1 field2=value2 ...>

    List rses.

    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = RSEClient(rucio_host=args.host, auth_host=args.auth_host,
                       account=args.account,
                       auth_type=args.auth_strategy, creds=creds,
                       ca_cert=args.ca_certificate, timeout=args.timeout)

    rse_expression = None
    if args.rse_expression:
        rse_expression = args.rse_expression

    try:
        rses = client.list_rses(rse_expression)
    except Exception, e:
        print 'Failed to list RSEs'
        print e
        return FAILURE

    for rse in rses:
        print '%(rse)s' % rse

    return SUCCESS


def list_rse_usage(args):
    """
    %(prog)s list-rse-usage [options] <rse>

    Show the space usage of a given rse

    """
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}
    else:
        creds = None

    rseclient = RSEClient(rucio_host=args.host, auth_host=args.auth_host,
                          account=args.account,
                          auth_type=args.auth_strategy, creds=creds,
                          ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        usages = rseclient.get_rse_usage(rse=args.rse)
    except Exception, e:
            print 'Failed to get RSE info'
            print e
            return FAILURE
    print 'Usage:'
    print '======'
    for usage in usages:
        for elem in usage:
            print '  ' + elem + ': ' + str(usage[elem])
        print '  ======'

    return SUCCESS


def list_account_limits(args):
    """
    %(prog)s list [options] <field1=value1 field2=value2 ...>

    List account limits.

    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = AccountClient(rucio_host=args.host, auth_host=args.auth_host,
                           account=args.account,
                           auth_type=args.auth_strategy, creds=creds,
                           ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        if args.rse:
            limits = client.get_account_limit(account=args.limit_account, rse=args.rse)
        else:
            limits = client.get_account_limits(account=args.limit_account)
    except Exception, e:
        print 'Failed to list Account limits'
        print e
        return FAILURE

    print "RSE: limit (bytes)"
    print "=================="
    for limit in limits.items():
        print "%s: %s" % limit

    return SUCCESS


def list_account_usage(args):
    """
    %(prog)s list [options] <field1=value1 field2=value2 ...>

    List account usage.

    """
    creds = None
    if args.auth_strategy == 'userpass':
        creds = {'username': args.username, 'password': args.password}

    client = AccountClient(rucio_host=args.host, auth_host=args.auth_host,
                           account=args.account,
                           auth_type=args.auth_strategy, creds=creds,
                           ca_cert=args.ca_certificate, timeout=args.timeout)

    try:
        if args.rse:
            usage = client.get_account_usage(account=args.usage_account, rse=args.rse)
        else:
            usage = client.get_account_usage(account=args.usage_account)
    except Exception, e:
        print 'Failed to list Account usage'
        print e
        return FAILURE

    print "RSE: usage (bytes), limit (bytes), quota_left (bytes)"
    print "====================================================="
    for item in usage:
        print "%s: %s, %s, %s" % (item['rse'], item['bytes'], item['bytes_limit'], item['bytes_remaining'])

    return SUCCESS


if __name__ == '__main__':
    usage = """
usage: %(prog)s <command> [options] [args]

Commands:

    help <command>  Output help for one of the commands below

"""
    oparser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]), add_help=True)
    subparsers = oparser.add_subparsers()

    # Main arguments
    oparser.add_argument('--version', action='version', version='%(prog)s ' + version.version_string())
    oparser.add_argument('--verbose', '-v', default=False, action='store_true', help="Print more verbose output")
    oparser.add_argument('-H', '--host', dest="host", metavar="ADDRESS", help="The Rucio API host")
    oparser.add_argument('--auth-host', dest="auth_host", metavar="ADDRESS", help="The Rucio Authentication host")
    oparser.add_argument('-a', '--account', dest="account", metavar="ACCOUNT", help="Rucio account to use")
    oparser.add_argument('-S', '--auth-strategy', dest="auth_strategy", default=None, help="Authentication strategy (userpass or x509 or ...)")
    oparser.add_argument('-T', '--timeout', dest="timeout", type=float, default=None, help="Set all timeout values to SECONDS")

    # Options for the userpass auth_strategy
    oparser.add_argument('-u', '--user', dest='username', default=None, help='username')
    oparser.add_argument('-pwd', '--password', dest='password', default=None, help='password')

    # Options for the x509  auth_strategy
    oparser.add_argument('--certificate', dest='certificate', default=None, help='Client certificate file')
    oparser.add_argument('--ca-certificate', dest='ca_certificate', default=None, help='CA certificate to verify peer against (SSL)')

    # Ping command
    ping_parser = subparsers.add_parser('ping', help='Ping Rucio server')
    ping_parser.set_defaults(which='ping')

    # The whoami command
    whoami_parser = subparsers.add_parser('whoami', help='Get information about account whose token is used')
    whoami_parser.set_defaults(which='whoami_account')

    # The list_replicas command
    list_replicas_parser = subparsers.add_parser('list-replicas', help='List file replicas')
    list_replicas_parser.set_defaults(which='list_replicas')
    list_replicas_parser.add_argument('--protocols', dest='protocols', action='store', help='List of comma separated protocols', required=False)
    list_replicas_parser.add_argument(dest='dids', nargs='+', action='store', help='List of space separated data identifiers')

    # The add-dataset command
    add_dataset_parser = subparsers.add_parser('add-dataset', help='Add dataset')
    add_dataset_parser.set_defaults(which='add_dataset')
    add_dataset_parser.add_argument('--monotonic', action='store_true', help='Monotonic status to True')
    add_dataset_parser.add_argument(dest='did', action='store', help='Dataset name to add')

    # The add-container command
    add_container_parser = subparsers.add_parser('add-container', help='Add container')
    add_container_parser.set_defaults(which='add_container')
    add_container_parser.add_argument('--monotonic', action='store_true', help='Monotonic status to True')
    add_container_parser.add_argument(dest='did', action='store', help='Container name to add')

    # TODO:
    # delete
    # regexp to validate a rse attribute key / value
    # add meta

    # The attach command
    attach_parser = subparsers.add_parser('attach', help='Attach a list of Data Identifiers (file, dataset or container) to an other Data Identifier (dataset or container)',
                                          description='Attach a list of Data Identifiers (file, dataset or container) to an other Data Identifier (dataset or container)')
    attach_parser.set_defaults(which='attach')
    attach_parser.add_argument(dest='dids', nargs='+', action='store', help='List of space separated data identifiers')
    attach_parser.add_argument('--to', dest='todid', default=None, action='store', help='Target data identifier', required=True)

    # The detach command
    detach_parser = subparsers.add_parser('detach', help='Detach a list of Data Identifiers (file, dataset or container) from an other Data Identifier (dataset or container)',
                                          description='Detach a list of Data Identifiers (file, dataset or container) from an other Data Identifier (dataset or container)')
    detach_parser.set_defaults(which='detach')
    detach_parser.add_argument(dest='dids', nargs='+', action='store', help='List of space separated data identifiers')
    detach_parser.add_argument('--from', dest='fromdid', default=None, action='store', help='Target data identifier', required=True)

    # The add_files_to_dataset command
    add_files_to_dataset_parser = subparsers.add_parser('add-files-to-dataset', help='Attach a list of files to a dataset', description='Attach a list of files to a dataset')
    add_files_to_dataset_parser.set_defaults(which='add-files-to-dataset')
    add_files_to_dataset_parser.add_argument(dest='files', nargs='+', action='store', help='List of space separated files')
    add_files_to_dataset_parser.add_argument('--to', dest='todid', default=None, action='store', help='Target dataset', required=True)

    # The add_datasets_to_container command
    add_datasets_to_container_parser = subparsers.add_parser('add-datasets-to-container', help='Attach a list of datasets to a container', description='Attach a list of datasets to a container')
    add_datasets_to_container_parser.set_defaults(which='add-datasets-to-container')
    add_datasets_to_container_parser.add_argument(dest='datasets', nargs='+', action='store', help='List of space separated datasets')
    add_datasets_to_container_parser.add_argument('--to', dest='todid', default=None, action='store', help='Target container', required=True)

    # The add_containers_to_container command
    add_containers_to_container_parser = subparsers.add_parser('add-containers-to-container', help='Attach a list of containers to a container', description='Attach a list of containers to a container')
    add_containers_to_container_parser.set_defaults(which='add-containers-to-container')
    add_containers_to_container_parser.add_argument(dest='containers', nargs='+', action='store', help='List of space separated datasets')
    add_containers_to_container_parser.add_argument('--to', dest='todid', default=None, action='store', help='Target container', required=True)

    # The list command
    list_parser = subparsers.add_parser('list-dids', help='List data identifier contents', description='List data identifier contents')
    list_parser.set_defaults(which='list-dids')
    list_parser.add_argument('--scope', dest='scope', action='store', default=None, help='Scope name')
    list_parser.add_argument('--recursive', dest='recursive', action='store_true', default=False, help='List data identifiers recursively')
    list_parser.add_argument('--flat', dest='flat', action='store_true', default=False, help='Flatl (unique) listing')
    list_parser.add_argument('--filter', dest='filter', action='store', help='Filter arguments in form `key=value another_key=next_value`')
    list_parser.add_argument(dest='dids', nargs='*', action='store', default=None, help='List of space separated data identifiers')

    # The list-scopes command
    scope_list_parser = subparsers.add_parser('list-scopes', help='List all available scopes')
    scope_list_parser.set_defaults(which='list_scopes')

    # The close command
    close_parser = subparsers.add_parser('close', help='Close data identifier')
    close_parser.set_defaults(which='close')
    close_parser.add_argument(dest='dids', nargs='+', action='store', help='List of space separated data identifiers')

    # The stat command
    stat_parser = subparsers.add_parser('stat', help='List attributes and statuses about data identifiers')
    stat_parser.set_defaults(which='stat')
    stat_parser.add_argument(dest='dids', nargs='+', action='store', help='List of space separated data identifiers')

    # The delete command
    delete_parser = subparsers.add_parser('delete', help='Delete data identifier')
    delete_parser.set_defaults(which='delete')
    delete_parser.add_argument(dest='dids', nargs='+', action='store', help='List of space separated data identifiers')
    delete_parser.add_argument('--from', dest='from', default=None, action='store', help='Target data identifier')

    # The list_files command
    list_files_parser = subparsers.add_parser('list-files', help='List data identifier contents')
    list_files_parser.set_defaults(which='list_files')
    list_files_parser.add_argument(dest='dids', nargs='+', action='store', help='List of space separated data identifiers')

    # The upload subparser
    upload_parser = subparsers.add_parser('upload', help='Upload method')
    upload_parser.set_defaults(which='upload')
    upload_parser.add_argument('--files', dest='files', nargs='+', action='store', help='List of space separated files', required=True)
    upload_parser.add_argument('--rse', dest='rse', action='store', help='Rucio Storage Element (RSE) name', required=True)
    upload_parser.add_argument('--scope', dest='scope', action='store', help='Scope name', required=True)  # to be optional after
    upload_parser.add_argument('--did', dest='did', action='store', help='Data identifier to add')
    upload_parser.add_argument('--project', dest='project', action='store', help='Project of the dataset')
    upload_parser.add_argument('--datatype', dest='datatype', action='store', help='Datatype of the dataset')
    upload_parser.add_argument('--run_number', dest='run_number', action='store', help='Run number of the dataset')
    upload_parser.add_argument('--stream_name', dest='stream_name', action='store', help='Stream name of the dataset')
    upload_parser.add_argument('--prod_step', dest='prod_step', action='store', help='Production step of the dataset')
    upload_parser.add_argument('--version', dest='version', action='store', help='Version of the dataset')
    upload_parser.add_argument('--campaign', dest='campaign', action='store', help='Campaign of the dataset')

    # The download subparser
    download_parser = subparsers.add_parser('download', help='Download method')
    download_parser.set_defaults(which='download')
    download_parser.add_argument('--dir', dest='dir', default='.', action='store', help='The directory to store the downloaded file.')
    download_parser.add_argument(dest='dids', nargs='+', action='store', help='List of space separated data identifiers')

    # The get-metadata subparser
    get_metadata_parser = subparsers.add_parser('get-metadata', help='get-metadata method')
    get_metadata_parser.set_defaults(which='get_metadata')
    get_metadata_parser.add_argument(dest='dids', nargs='+', action='store', help='List of space separated data identifiers')

    # The set-metadata subparser
    set_metadata_parser = subparsers.add_parser('set-metadata', help='set-metadata method')
    set_metadata_parser.set_defaults(which='set_metadata')
    set_metadata_parser.add_argument('--scope', dest='scope', action='store', help='Scope name', required=True)  # to be optional after
    set_metadata_parser.add_argument('--did', dest='did', action='store', help='Data identifier to add', required=True)
    set_metadata_parser.add_argument('--key', dest='key', action='store', help='Attribute key', required=True)
    set_metadata_parser.add_argument('--value', dest='value', action='store', help='Attribute value', required=True)

    # The delete-metadata subparser
    delete_metadata_parser = subparsers.add_parser('delete-metadata', help='delete-metadata method')
    delete_metadata_parser.set_defaults(which='delete_metadata')

    # The list-rse-usage subparser
    list_rse_usage_parser = subparsers.add_parser('list-rse-usage', help='list-rse-usage method')
    list_rse_usage_parser.set_defaults(which='list_rse_usage')
    list_rse_usage_parser.add_argument(dest='rse', action='store', help='Rucio Storage Element (RSE) name')
    list_rse_usage_parser.add_argument('--history', dest='history', default=False, action='store', help='List rse usage history')

    # The list-account-usage subparser
    list_account_usage_parser = subparsers.add_parser('list-account-usage', help='list-account-usage method')
    list_account_usage_parser.set_defaults(which='list_account_usage')
    list_account_usage_parser.add_argument(dest='usage_account', action='store', help='Account name')
    list_account_usage_parser.add_argument('--rse', action='store', help='Show usage for this rse')

    # The list-account-limits subparser
    list_account_limits_parser = subparsers.add_parser('list-account-limits', help='List account limits on RSEs')
    list_account_limits_parser.set_defaults(which='list_account_limits')
    list_account_limits_parser.add_argument('limit_account', action='store', help='Account name')
    list_account_limits_parser.add_argument('--rse', dest='rse', action='store', help='RSE name')

    # Add replication rule subparser
    add_rule_parser = subparsers.add_parser('add-rule', help='Add replication rule')
    add_rule_parser.set_defaults(which='add_rule')
    add_rule_parser.add_argument(dest='dids', action='store', nargs='+', help='DID(s) to apply the rule to')
    add_rule_parser.add_argument(dest='copies', action='store', type=int, help='Number of copies')
    add_rule_parser.add_argument(dest='rse_expression', action='store', help='RSE Expression')
    add_rule_parser.add_argument('--weight', dest='weight', action='store', help='RSE Weight')
    add_rule_parser.add_argument('--lifetime', dest='lifetime', action='store', type=int, help='Rule lifetime (in seconds)')
    add_rule_parser.add_argument('--grouping', dest='grouping', action='store', choices=['DATASET', 'ALL', 'NONE'], help='Rule grouping')
    add_rule_parser.add_argument('--locked', dest='locked', action='store_true', help='Rule locking')
    add_rule_parser.add_argument('--source-replice-expression', dest='source_replica_expression', action='store', help='Source Replica Expression')
    add_rule_parser.add_argument('--notify', dest='notify', action='store', help='Notification strategy : Y (Yes), N (No), C (Close)')

    # Delete replication rule subparser
    delete_rule_parser = subparsers.add_parser('delete-rule', help='Delete replication rule')
    delete_rule_parser.set_defaults(which='delete_rule')
    delete_rule_parser.add_argument(dest='rule_id', action='store', help='Rule id')

    # Info replication rule subparser
    info_rule_parser = subparsers.add_parser('rule-info', help='Retrieve information about a rule')
    info_rule_parser.set_defaults(which='info_rule')
    info_rule_parser.add_argument(dest='rule_id', action='store', help='Rule id')

    # The list_rules command
    list_rules_parser = subparsers.add_parser('list-rules', help='List replication rules')
    list_rules_parser.set_defaults(which='list_rules')
    list_rules_parser.add_argument('--rule', dest='rule_id', action='store', help='List by rule id')
    list_rules_parser.add_argument('--did', dest='did', action='store', help='List by did')
    list_rules_parser.add_argument('--file', dest='file', action='store', help='List associated rules of an affected file')
    list_rules_parser.add_argument('--account', dest='rule_account', action='store', help='List by account')
    list_rules_parser.add_argument('--subscription', dest='subscription', action='store', help='List by account and subscription name', metavar=('ACCOUNT', 'SUBSCRIPTION'), nargs=2)

    # The update_rule command
    update_rule_parser = subparsers.add_parser('update-rule', help='Update replication rule')
    update_rule_parser.set_defaults(which='update_rule')
    update_rule_parser.add_argument(dest='rule_id', action='store', help='Rule id')
    update_rule_parser.add_argument('--lifetime', dest='lifetime', action='store', help='Lifetime in seconds.')
    update_rule_parser.add_argument('--locked', dest='locked', action='store', help='Locked (True/False).')

    # The list-rses command
    list_rses_parser = subparsers.add_parser('list-rses', help='List RSEs')
    list_rses_parser.set_defaults(which='list_rses')
    list_rses_parser.add_argument('--expression', dest='rse_expression', action='store', help='RSE Expression')

    argcomplete.autocomplete(oparser)

    if len(sys.argv) == 1:
        oparser.print_help()
        sys.exit(FAILURE)

    args = oparser.parse_args(sys.argv[1:])

    commands = {'ping': ping,
                'list_replicas': list_replicas,
                'attach': attach,
                'detach': detach,
                'add-files-to-dataset': add_files_to_dataset,
                'add-datasets-to-container': add_datasets_to_container,
                'add-containers-to-container': add_containers_to_container,
                'add_dataset': add_dataset,
                'add_container': add_container,
                'delete': delete,
                'list-dids': list_dids,
                'list_files': list_files,
                'list_scopes': list_scopes,
                'close': close,
                'stat': stat,
                'whoami_account': whoami_account,
                'download': download,
                'upload': upload,
                'get_metadata': get_metadata,
                'set_metadata': set_metadata,
                'delete_metadata': delete_metadata,
                'add_rule': add_rule,
                'info_rule': info_rule,
                'list_rules': list_rules,
                'delete_rule': delete_rule,
                'list_rses': list_rses,
                'list_rse_usage': list_rse_usage,
                'list_account_limits': list_account_limits,
                'list_account_usage': list_account_usage,
                'update_rule': update_rule}

    try:
        start_time = time.time()
        command = commands.get(args.which)
        result = command(args)
        end_time = time.time()
        if args.verbose:
            print "Completed in %-0.4f sec." % (end_time - start_time)
        sys.exit(result)
    except (RuntimeError, NotImplementedError), e:
        print >> sys.stderr, "ERROR: ", e
        sys.exit(FAILURE)
