#!/usr/bin/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>, 2016-2017

"""
Get the average megabytes per second per file on a link from the dashboard.
"""

import datetime
import json
import requests


def get_link_throughput(minutes):

    """ Retrieve the data from the dashboard. """

    url = 'http://dashb-atlas-ddm.cern.ch/dashboard/request.py/matrix.json'\
          '?activity=Analysis+Input&activity=Data+Brokering&activity=Data+Consolidation'\
          '&activity=Data+Rebalancing&activity=Deletion&activity=Express&activity=Functional+Test'\
          '&activity=Production+Input&activity=Production+Output&activity=Recovery'\
          '&activity=SFO+to+EOS+export&activity=Staging&activity=T0+Export&activity=T0+Tape'\
          '&activity=User+Subscriptions&activity=default'\
          '&src_grouping=site&src_grouping=token&dst_grouping=site&dst_grouping=token&interval=' + str(minutes)

    res = requests.get(url).json()['transfers']['rows']
    return res


if __name__ == '__main__':

    DATA = {}

    THROUGHPUTS = [('1h', get_link_throughput(60)),
                   ('1d', get_link_throughput(1440)),
                   ('1w', get_link_throughput(10080))]

    for metric, throughput in THROUGHPUTS:

        for row in throughput:
            src = '%s_%s' % (row[0], row[1])
            dst = '%s_%s' % (row[2], row[3])
            link = '%s:%s' % (src, dst)
            if link.startswith(':') or link.endswith(':') or link.endswith('_None') or '_None:' in link:
                continue

            if link not in DATA.keys():
                DATA[link] = {'mbps': {'dashb': {'1h': [],
                                                 '1d': [],
                                                 '1w': [],
                                                 'timestamp': datetime.datetime.utcnow().strftime('%Y-%m-%dT%H:%M:%S')}}}

            if metric == '1h':
                if row[4] != 0:
                    DATA[link]['mbps']['dashb']['1h'] = row[4]
                    DATA[link]['mbps']['dashb']['1d'] = row[4]
                    DATA[link]['mbps']['dashb']['1w'] = row[4]

            if metric == '1d':
                if row[4] != 0:
                    DATA[link]['mbps']['dashb']['1d'] = row[4]
                    DATA[link]['mbps']['dashb']['1w'] = row[4]

            if metric == '1w':
                if row[4] != 0:
                    DATA[link]['mbps']['dashb']['1w'] = row[4]

    MAPPING = None
    with open('/data/metrix/data/mapping-rse-site/latest.json', 'r') as f:
        MAPPING = json.load(f)

    for link in DATA:
        src, dst = link.split(':')
        DATA[link]['src_site'] = MAPPING[src]
        DATA[link]['dst_site'] = MAPPING[dst]

    SITE_DATA = {}
    for link in DATA:
        site_link = '%s:%s' % (DATA[link]['src_site'], DATA[link]['dst_site'])
        if site_link in SITE_DATA.keys():
            if DATA[link]['mbps']['dashb']['1h'] != []:
                if DATA[link]['mbps']['dashb']['1h'] != 0:
                    SITE_DATA[site_link]['mbps']['dashb']['1h'].append(DATA[link]['mbps']['dashb']['1h'])
            if DATA[link]['mbps']['dashb']['1d'] != []:
                if DATA[link]['mbps']['dashb']['1d'] != 0:
                    SITE_DATA[site_link]['mbps']['dashb']['1d'].append(DATA[link]['mbps']['dashb']['1d'])
            if DATA[link]['mbps']['dashb']['1w'] != []:
                if DATA[link]['mbps']['dashb']['1w'] != 0:
                    SITE_DATA[site_link]['mbps']['dashb']['1w'].append(DATA[link]['mbps']['dashb']['1w'])
        else:
            SITE_DATA[site_link] = {'mbps': DATA[link]['mbps']}
            SITE_DATA[site_link]['mbps']['dashb']['1h'] = [SITE_DATA[site_link]['mbps']['dashb']['1h']]
            SITE_DATA[site_link]['mbps']['dashb']['1d'] = [SITE_DATA[site_link]['mbps']['dashb']['1d']]
            SITE_DATA[site_link]['mbps']['dashb']['1w'] = [SITE_DATA[site_link]['mbps']['dashb']['1w']]

    for link in SITE_DATA:
        if SITE_DATA[link]['mbps']['dashb']['1h'] == [[]]:
            SITE_DATA[link]['mbps']['dashb']['1h'] = [0.0]
        if SITE_DATA[link]['mbps']['dashb']['1d'] == [[]]:
            SITE_DATA[link]['mbps']['dashb']['1d'] = [0.0]
        if SITE_DATA[link]['mbps']['dashb']['1w'] == [[]]:
            SITE_DATA[link]['mbps']['dashb']['1w'] = [0.0]

    for link in SITE_DATA:
        if [] in SITE_DATA[link]['mbps']['dashb']['1h']:
            SITE_DATA[link]['mbps']['dashb']['1h'].remove([])
        if [] in SITE_DATA[link]['mbps']['dashb']['1d']:
            SITE_DATA[link]['mbps']['dashb']['1d'].remove([])
        if [] in SITE_DATA[link]['mbps']['dashb']['1w']:
            SITE_DATA[link]['mbps']['dashb']['1w'].remove([])

    for link in SITE_DATA:
        SITE_DATA[link]['mbps']['dashb']['1h'] = round(sum(SITE_DATA[link]['mbps']['dashb']['1h']) * 0.0000000002778, 2)
        SITE_DATA[link]['mbps']['dashb']['1d'] = round(sum(SITE_DATA[link]['mbps']['dashb']['1d']) * 0.00000000001157, 2)
        SITE_DATA[link]['mbps']['dashb']['1w'] = round(sum(SITE_DATA[link]['mbps']['dashb']['1w']) * 0.000000000001653, 2)

    for link in SITE_DATA:
        if SITE_DATA[link]['mbps']['dashb']['1h'] == 0:
            del SITE_DATA[link]['mbps']['dashb']['1h']
        if SITE_DATA[link]['mbps']['dashb']['1d'] == 0:
            del SITE_DATA[link]['mbps']['dashb']['1d']
        if SITE_DATA[link]['mbps']['dashb']['1w'] == 0:
            del SITE_DATA[link]['mbps']['dashb']['1w']

    with open('/data/metrix/data/mbps-dashb/mbps-dashb-{0}.json'.format(datetime.datetime.utcnow().isoformat()[:-7]), 'w') as f:
        json.dump(SITE_DATA, f, indent=1, sort_keys=True)

    with open('/data/metrix/data/mbps-dashb/latest.json', 'w') as f:
        json.dump(SITE_DATA, f, indent=1, sort_keys=True)
