#!/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:
# - Joaquin Bogado, <joaquin.bogado@cern.ch>, 2014

import sys
import time
import datetime
import requests
import socket
from rucio.client import Client
from rucio.common.utils import uuid

c = Client()
scp = ''
fl = []
dsname = 'tests.rucio_client_test_bm_' + uuid().urn.split(':')[2]


def test_create_dataset(file_list, dsn=dsname):
    c.add_dataset(scp, dsn)
    tick = time.time()
    c.add_files_to_dataset(scp, dsn, file_list[0:100])
    tick = time.time() - tick
    c.add_files_to_dataset(scp, dsn, file_list[100:600])
    c.add_files_to_dataset(scp, dsn, file_list[600:900])
    c.add_files_to_dataset(scp, dsn, file_list[900:1200])
    return tick


def test_read_big_dataset():
    tick = time.time()
    c.list_files(scp, dsname)
    return time.time() - tick


def test_quering_replicas_chuncked(file_list):
    tick = time.time()
    c.list_replicas(file_list[0:100])
    c.list_replicas(file_list[100:200])
    c.list_replicas(file_list[200:300])
    c.list_replicas(file_list[300:400])
    c.list_replicas(file_list[400:500])
    c.list_replicas(file_list[500:600])
    c.list_replicas(file_list[600:700])
    c.list_replicas(file_list[700:800])
    c.list_replicas(file_list[800:900])
    c.list_replicas(file_list[900:1000])
    return (time.time() - tick)/10


def test_create_rule():
    tick = time.time()
    c.add_replication_rule([{'scope': scp, 'name': dsname}], 1, 'MOCK', lifetime=3600)   # lifetime = 1 hour. Should not trigger replications because all the files are in MOCK rse
    return time.time() - tick


def get_filelist():
    # get the offset for the subset of files.
    try:
        f = open('/tmp/.slice')
        s = int(f.read())
        f.close()
    except IOError:
        s = 0
    f = open('/tmp/.slice', 'w')
    f.write(str((s+1) % 20))   # the offset is a number between 0 and 19
    f.close()
    files = ('rucio_test_mock_file_{0:05}'.format(x) for x in xrange(20000))
    f = []
    f.extend(files)
    fl = []
    for i in (({'scope': 'mock', 'name': x}) for x in f[(s*1000): (s*1000)+1000]):
        fl.append(i)
    return fl


def main():
    global scp
    import argparse
    parser = argparse.ArgumentParser(description='Rucio client test')
    parser.add_argument('-r', '--report', dest='report', default=False, action='store_true', help='Report the results to xsls.cern.ch')
    parser.add_argument('-s', '--scope', dest='scp', default='', help='Scope for testing. Usually user.<account_name>', required=True)
    parser.add_argument('-q', '--quiet', dest='quiet', default=False, action='store_true', help='No xml output')
    args = parser.parse_args()

    scp = args.scp

    fl = get_filelist()

    timestamp = datetime.datetime.now().isoformat().split('.')[0]

    createtime = test_create_dataset(fl)
    readtime = test_read_big_dataset()
    querytime = test_quering_replicas_chuncked(fl)
    ruletime = test_create_rule()

    xmlmetric = """<?xml version="1.0" encoding="UTF-8"?>
<serviceupdate xmlns="http://sls.cern.ch/SLS/XML/update">
<id>rucio.{5}.client_benchmark</id>
<timestamp>{0}</timestamp>
<availability>100</availability>
<data>
<numericvalue name="rucio.{5}.client_benchmark.create_dataset" desc="Time to create a dataset with 100 files in seconds">{2}</numericvalue>
<numericvalue name="rucio.{5}.client_benchmark.read_big_dataset" desc="Time to read dataset with 1000 files in seconds">{1}</numericvalue>
<numericvalue name="rucio.{5}.client_benchmark.quering_replicas_chuncked" desc="Avg time to query the replicas of 100 files in seconds">{3}</numericvalue>
<numericvalue name="rucio.{5}.client_benchmark.add_rule" desc="Time to add a rule to a dataset in seconds">{4}</numericvalue>
</data>
</serviceupdate>""".format(timestamp, readtime, createtime, querytime, ruletime, socket.getfqdn()[:-8])
    if not args.quiet:
        print xmlmetric
    if args.report is True:
        if not args.quiet:
            print "Reporting to http://xsls.cern.ch"
        r = requests.post('http://xsls.cern.ch', files={'file': xmlmetric})
        sys.exit(not(r.status_code == 200))

if __name__ == '__main__':
    main()
