#!/usr/bin/env python
"""
Benchmarking tool for scan-http-service.

Requires some Python libraries that are not included in the standard Python
installation:
 - requests.
"""
from __future__ import print_function

import requests
import datetime
import multiprocessing
import sys


def _dt_to_js_iso8601(dt):
    '''
    Print datetimes in a format like 2012-04-23T18:25:43.511Z.

    We want everything to be explicitly in UTC.
    '''
    if dt.tzinfo is not None and dt.utcoffset() != datetime.timedelta(0):
        raise Exception("Only UTC (Z) datetimes supported")

    without_micros_part = u'{:%Y-%m-%dT%H:%M:%S}'.format(dt)
    micros_part = u'{:%f}'.format(dt).rstrip('0')
    if len(micros_part) == 0:
        micros_part = '0'
    return u'{0}.{1}Z'.format(without_micros_part, micros_part)

def dbg(*args, **kwargs):
    from inspect import stack
    from sys import stderr
    "Convenience function for development time print-line debugging."
    calling_function = stack()[1][3]
    kwargs['file'] = stderr
    return print('#', calling_function, *args, **kwargs)

def timed_call(func, args, kwargs={}):
    """Returns a two-tuple of time_in_seconds, result."""
    start = datetime.datetime.now()
    result = func(*args, **kwargs)
    end = datetime.datetime.now()
    return (end - start).total_seconds(), result

def timed_get(mime_type_url_pair):
    dbg('GET', mime_type_url_pair)
    mime_type, url = mime_type_url_pair
    return timed_call(requests.get, [url], kwargs={'headers': {"Accept": mime_type}})

def bench(base_url, hours, process_count, media_type):
    def get_scans(start=None):
        if start is None:
            url = u'{}scans/'.format(base_url)
        else:
            url = u'{0}scans/?start={1}'.format(base_url, _dt_to_js_iso8601(start))

        seconds, r = timed_call(requests.get, [url])
        dbg(r.status_code)
        assert r.status_code == 200, r
        dbg('Content-Type:', r.headers['content-type'])
        dbg('Encoding:', r.encoding)

        deserialized = r.json()
        scans = deserialized['scans']
        dbg('Scan count:', len(scans))
        return scans

    start = datetime.datetime.now() - datetime.timedelta(hours=hours)
    scans = get_scans(start=start)

    benchmark_urls = [] # will contain pairs like (mime_type, url)

    for scan in scans:
        dbg('Scan (task name, time, id):', scan['taskName'], scan['timestamp'], scan['id'])
        for moment_ref in scan['moments']:
            label = moment_ref['moment']['label']
            url = u'{}{}'.format(base_url, moment_ref['href'])
            benchmark_urls.append((media_type, url))

    pool = multiprocessing.Pool(process_count)
    gets_started = datetime.datetime.now()
    result = pool.map_async(timed_get, benchmark_urls).get(99999999999)
    gets_ended = datetime.datetime.now()
    times = [pair[0] for pair in result]
    count = len(times)

    if len(times) == 0:
        dbg("No times. Add --hours argument with a higher value?")
        sys.exit(0)

    time_sum = reduce(lambda x, y: x+y, times)
    times_sorted = sorted(times)

    dbg("Moments total:", count)
    dbg("Took (incl scheduling):", (gets_ended - gets_started).total_seconds())
    dbg("Took seconds (only GET):", time_sum)
    dbg("Moments per sec:", (gets_ended - gets_started).total_seconds() / float(count))
    dbg("Median time:", times_sorted[len(times_sorted) / 2])
    dbg("Mean time:", float(time_sum) / count)
    dbg("Min:", times_sorted[0])
    dbg("Max:", times_sorted[-1])


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--base-url', dest='base_url', required=True,
                        help="Base URL for scan service, like http://localhost:34280/api/v0/")
    parser.add_argument('--hours', dest='hours', default=48,
                        type=int)
    parser.add_argument('--process-count', dest='process_count', default=1,
                        type=int)
    parser.add_argument('--media-type', dest='media_type',
                        choices=['application/javascript', 'application/vnd.vaisala.iris.protobuf0'],
                        default='application/javascript')

    args = parser.parse_args()
    bench(args.base_url, args.hours, args.process_count, args.media_type)
