#!/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
# - Hannes Hansen, <hannes.jakob.hansen@cern.ch>, 2019

"""
Get the average megabytes per second per file on a link from ElasticSearch/FTS.
"""

import datetime
import json
import time

from elasticsearch import Elasticsearch
from six import iteritems


def dict_merge(dct, merge_dct):

    """ Merge two dictionaries. """

    for k, v in iteritems(merge_dct):  # pylint: disable=invalid-name,unused-variable
        if k in dct and isinstance(dct[k], dict and isinstance(merge_dct[k], dict)):
            dict_merge(dct[k], merge_dct[k])
        else:
            dct[k] = merge_dct[k]


def get_mbps(label, hours):

    """ Retrieve the data from ElasticSearch. """

    begin_time = int(time.mktime((datetime.datetime.now() - datetime.timedelta(hours=hours)).timetuple()) * 1000)
    end_time = int(time.mktime((datetime.datetime.now()).timetuple()) * 1000)

    body = """
{
  "size": 0,
  "query": {
    "bool": {
      "must": [
        {
          "query_string": {
            "query": "_type:transfer-done",
            "analyze_wildcard": true
          }
        },
        {
          "range": {
            "@timestamp": {
              "gte": %s,
              "lte": %s,
              "format": "epoch_millis"
            }
          }
        }
      ],
      "must_not": []
    }
  },
  "_source": {
    "excludes": []
  },
  "aggs": {
    "2": {
      "terms": {
        "field": "payload.src-rse",
        "size": 999,
        "order": {
          "_term": "desc"
        }
      },
      "aggs": {
        "3": {
          "terms": {
            "field": "payload.dst-rse",
            "size": 999,
            "order": {
              "_term": "desc"
            }
          },
          "aggs": {
            "1": {
              "percentiles": {
                "field": "payload.duration",
                "percents": [
                  95
                ],
                "keyed": false
              }
            }
          }
        }
      }
    }
  }
}
""" % (begin_time, end_time)

    esearch = Elasticsearch([{'host': 'atlas-kibana.mwt2.org', 'port': 9200}])
    res = esearch.search(index='rucio-events-2017*', body=body)

    data = {}
    for src_site in res['aggregations']['2']['buckets']:
        for dst_site in src_site['3']['buckets']:
            tmp_link = '%s:%s' % (src_site['key'], dst_site['key'])
            if tmp_link.startswith(':') or tmp_link.endswith(':'):
                continue
            throughput = dst_site['1']['values'][0]['value']
            if throughput == 'Infinity':
                pass
            else:
                DATA[tmp_link] = {'mbps': {'fts': {label: round(throughput, 2)}}}

    return data


if __name__ == '__main__':

    DATA = get_mbps('1h', 1)
    DAYS = get_mbps('1d', 24)
    WEEKS = get_mbps('1w', 168)

    dict_merge(DATA, DAYS)
    dict_merge(DATA, WEEKS)

    for link in DATA:
        if len(DATA[link]['mbps']['fts'].keys()) != 3:
            if '1d' not in DATA[link]['mbps']['fts'].keys():
                if '1h' in DATA[link]['mbps']['fts'].keys():
                    DATA[link]['mbps']['fts']['1d'] = DATA[link]['mbps']['fts']['1h']
            if '1w' not in DATA[link]['mbps']['fts'].keys():
                DATA[link]['mbps']['fts']['1w'] = DATA[link]['mbps']['fts']['1d']
        DATA[link]['mbps']['fts']['timestamp'] = datetime.datetime.utcnow().isoformat()[:-7]

    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 '1h' in DATA[link]['mbps']['fts'].keys() and '1h' in SITE_DATA[site_link]['mbps']['fts'].keys():
                if DATA[link]['mbps']['fts']['1h'] > SITE_DATA[site_link]['mbps']['fts']['1h']:
                    SITE_DATA[site_link]['mbps']['fts']['1h'] = DATA[link]['mbps']['fts']['1h']
            if '1d' in DATA[link]['mbps']['fts'].keys() and '1d' in SITE_DATA[site_link]['mbps']['fts'].keys():
                if DATA[link]['mbps']['fts']['1d'] > SITE_DATA[site_link]['mbps']['fts']['1d']:
                    SITE_DATA[site_link]['mbps']['fts']['1d'] = DATA[link]['mbps']['fts']['1d']
            if '1w' in DATA[link]['mbps']['fts'].keys() and '1w' in SITE_DATA[site_link]['mbps']['fts'].keys():
                if DATA[link]['mbps']['fts']['1w'] > SITE_DATA[site_link]['mbps']['fts']['1w']:
                    SITE_DATA[site_link]['mbps']['fts']['1w'] = DATA[link]['mbps']['fts']['1w']
        else:
            SITE_DATA[site_link] = {'mbps': DATA[link]['mbps']}

    with open('/data/metrix/data/mbps-fts/mbps-fts-{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-fts/latest.json', 'w') as f:
        json.dump(SITE_DATA, f, indent=1, sort_keys=True)
