#!/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-2013
# - 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

"""
    Rucio CLI.
"""

import argcomplete
import argparse
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, FileReplicaAlreadyExists
from rucio.common.utils import adler32
from rucio.rse import rsemanager as rsemgr

SUCCESS = 0
FAILURE = 1

DEFAULT_SECURE_PORT = 443
DEFAULT_PORT = 80


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:
            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()) * ' '
    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:
        print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.todid, e)
        return FAILURE
    except Exception, e:
        print e

    try:
        client.attach_dids(scope=scope, name=name, dids=dids)
    except Exception, e:
        print 'Failed to attach data identifier'
        print 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:
        print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.fromdid, e)
        return FAILURE
    except Exception, e:
        print e

    try:
        client.detach_dids(scope=scope, name=name, dids=dids)
    except Exception, e:
        print 'Failed to detach data identifier'
        print 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


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}

    dsn = None
    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)
    if args.did:
        try:
            dsn = {'scope': args.did.split(':')[0], 'name': args.did.split(':')[1]}
        except ValueError, e:
            print 'ERROR cannot extract the scope and name from %s : [%s]' % (args.did, e)
            return FAILURE

    files = args.files
    list_files = []
    lfns = {}
    revert_dict = {}
    for name in files:
        size = os.stat(name).st_size
        checksum = adler32(name)
        list_files.append({'scope': args.scope, 'name': os.path.basename(name), 'bytes': size, 'adler32': checksum})
        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})
        revert_dict[args.scope, os.path.basename(name)] = os.path.dirname(name)

    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)
    add_files_individually = False
    if dsn:
        try:
            if args.account is None:
                account = client.whoami()['account']
            else:
                account = args.account
            try:
                client.add_dataset(scope=dsn['scope'], name=dsn['name'], rules=[{'account': account, 'copies': 1, 'rse_expression': args.rse, 'grouping': 'DATASET'}], meta=metadata)
            except DataIdentifierAlreadyExists, e:
                print e
            client.add_files_to_dataset(scope=dsn['scope'], name=dsn['name'], files=list_files, rse=args.rse)
            for dir in lfns:
                try:
                    rsemgr.upload(rse_settings=rse_settings, lfns=lfns[dir], source_dir=dir)
                    print 'Upload operation for %s:%s done' % (dsn['scope'], dsn['name'])
                except FileReplicaAlreadyExists, e:
                    add_files_individually = True
        except Duplicate, e:
            add_files_individually = True
        except FileAlreadyExists, e:
            add_files_individually = True

        if add_files_individually:
            print 'Will check all individual files'
            for did in list_files:
                add_files_individually = False
                try:
                    client.add_files_to_dataset(scope=dsn['scope'], name=dsn['name'], files=[did], rse=args.rse)
                except Duplicate, e:
                    print e
                    add_files_individually = True
                except FileAlreadyExists, e:
                    print e
                    add_files_individually = True
                if add_files_individually:
                    meta = client.get_metadata(did['scope'], did['name'])
                    if rsemgr.exists(rse_settings=rse_settings, files={'name': did['name'], 'scope': did['scope']}):
                        print 'File %s:%s already exists on RSE. Will not try to reupload' % (did['scope'], did['name'])
                        print 'TODO : checksum validation'
                    else:
                        if meta['adler32'] == did['adler32']:
                            print 'Local files and file recorded in Rucio have the same checksum. Will try the upload'
                            directory = revert_dict[did['scope'], did['name']]
                            rsemgr.upload(rse_settings=rse_settings, lfns=[{'name': did['name'], 'scope': did['scope']}], source_dir=directory)
                        else:
                            print 'Checksum mismatch between local (%s) and recorded file (%s)' % (meta['adler32'], did['adler32'])
    else:
        try:
            client.add_replicas(files=list_files, rse=args.rse)
            client.add_replication_rule(list_files, copies=1, rse_expression=args.rse)
            for dir in lfns:
                try:
                    rsemgr.upload(rse_settings=rse_settings, lfns=lfns[dir], source_dir=dir)
                    print 'Upload operation for %(files)s done' % locals()
                except FileReplicaAlreadyExists, e:
                    add_files_individually = True
        except Duplicate, e:
            add_files_individually = True

        if add_files_individually:
            print 'Will check all individual files'
            for did in list_files:
                add_files_individually = False
                try:
                    client.add_replicas(files=[did], rse=args.rse)
                except Duplicate, e:
                    print e
                    add_files_individually = True
                except FileAlreadyExists, e:
                    print e
                    add_files_individually = True
                if add_files_individually:
                    meta = client.get_metadata(did['scope'], did['name'])
                    if rsemgr.exists(rse_settings=rse_settings, files={'name': did['name'], 'scope': did['scope']}):
                        print 'File %s:%s already exists on RSE. Will not try to reupload' % (did['scope'], did['name'])
                        print 'TODO : checksum validation'
                    else:
                        if meta['adler32'] == did['adler32']:
                            print 'Local files and file recorded in Rucio have the same checksum. Will try the upload'
                            directory = revert_dict[did['scope'], did['name']]
                            rsemgr.upload(rse_settings=rse_settings, lfns=[{'name': did['name'], 'scope': did['scope']}], source_dir=directory)
                        else:
                            print 'Checksum mismatch between local (%s) and recorded file (%s)' % (meta['adler32'], did['adler32'])
        #except Exception, e:
        #    print 'Failed to upload %(files)s' % locals()
        #    print e


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 = {}
    for did in args.dids:
        try:
            scope, name = did.split(':')
            files = client.list_replicas([{'scope': scope, 'name': name}])
            for file in files:
                rse = random.choice(file['rses'].keys())
                if not rse in rse_dict:
                    rse_dict[rse] = rsemgr.get_rse_info(rse)
                rsemgr.download(rse_dict[rse], files=[{'name': file['name'], 'scope': file['scope'], 'adler32': file['adler32']}, ], dest_dir=args.dir, printstatements=True)
            print 'download operation for %s:%s done' % (scope, name)
        except ValueError, e:
            print 'ERROR cannot extract the scope and name from %s : [%s]' % (did, e)
            return FAILURE
        except Exception, e:
            print 'Failed to download %(scope)s:%(name)s' % locals()
            print e


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)
        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 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}

    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:
        raise NotImplementedError  # TODO
    except Exception, e:
        print 'Failed to delete 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:
            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(':')
            rules = client.list_did_rules(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.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:
            print 'Failed to list rules'
            print e
            return FAILURE
    else:
        print 'At least one option has to be given. Use -h to list the options.'
        return FAILURE

    for rule in rules:
        print rule

    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:
        usage = rseclient.get_rse_usage(rse=args.rse)
    except Exception, e:
            print 'Failed to get RSE info'
            print e
            return FAILURE
    print 'Usage:'
    print '======'
    for elem in usage:
        print '  ' + elem + ': ' + str(usage[elem])

    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 data identifiers in a scope')
    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='account', action='store', help='Account name')
    list_account_usage_parser.add_argument('--history', dest='history', default=False, action='store', help='List account usage history')

    # The list-account-limits subparser
    list_account_limits_parser = subparsers.add_parser('list-account-limits', help='list-account-limits method')
    list_account_limits_parser.set_defaults(which='list_account_limits')
    list_account_limits_parser.add_argument(dest='account', action='store', help='Account 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 hours)')
    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')

    # 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('--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 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': add,
                'add_dataset': add_dataset,
                'add_container': add_container,
                'delete': delete,
                # 'delete_dataset': delete_dataset,
                # 'delete_container': delete_container,
                '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}

    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)
