from __future__ import division

import os
import random
import time

from multiprocessing import Process, Queue

import datetime
from pymongo import MongoClient

from torii import Torii



class MonitoringBenchmark(object):

    def __init__(self,
                 mongo_url='mongodb://localhost:27017',
                 torii_url='http://localhost:8443/torii',
                 timeout=60, page_size=20,
                 nb_retry=3,
                 output_file=None,
                 ref_data_keys=[]):

        self.mongo_url = mongo_url
        self.torii_url = torii_url
        self.timeout = timeout
        self.page_size = page_size
        self.nb_run = 0
        self.nb_retry = nb_retry
        self.ref_data_keys = ref_data_keys

        if output_file:
            self.output_file = output_file
        else:
            self.output_file = 'torii_benchmark_{}_monitor.csv'.format(datetime.datetime.today().strftime('%Y%m%d-%H%M'))

        # Initiate several sessions
        # we don't care about different users since sshd is not involved in monitoring
        self.t = Torii(base_url=torii_url, username='hpcgadmin', password='hpcgadmin')

        self.logger = self.t.logger
        self.logger.setLevel('INFO')

        self.csv_file = open(self.output_file, 'w')
        self.logger.info('initiate output file {}'.format(os.path.abspath(self.output_file)))
        self.csv_file.write('##\n'
                            '## Torii Benchmark \n'
                            '## {0}\n'
                            '##\n\n\n'.format(datetime.datetime.today().strftime('%Y-%m-%d %H:%M:%S')))

        #                   ' {:10d} | {:50s} | {:10d} | {:10d} | {:10.1f} | {:10.1f} | {:30s} |'
        self.csv_file.write(' num        |'
                            ' use case                                           |'
                            ' method                         |'
                            ' nb procs   |')

        for key in self.ref_data_keys:
            self.csv_file.write(' {:20s} |'.format(key))

        self.csv_file.write(' nb reqs    |'
                            ' nb errors  |'
                            ' time(s)    |'
                            ' req/s      |'
                            ' timestamp  |'
                            ' date                           |\n')

        self.csv_file.flush()


    def access_task_page_sort(self, q, i, max_page=100):
        self.access_task_page(q, i, max_page=max_page, sorted=True, count=False, raw=False)

    def access_task_page_raw(self, q, i, max_page=100):
        self.access_task_page(q, i, max_page=max_page, sorted=True, count=False, raw=True)

    def access_task_page_count(self, q, i, max_page=100):
        self.access_task_page(q, i, max_page=max_page, sorted=True, count=True, raw=False)

    def access_task_page(self, q, i, max_page=100, sorted=False, count=False, raw=False):

        nb_accesses = 0
        nb_errors = 0

        t_local = Torii(base_url=self.torii_url, username='hpcgadmin', password='hpcgadmin')

        nb_tasks = t_local.tasks.count()
        nb_pages = nb_tasks // self.page_size

        walltime = time.time() + self.timeout

        self.logger.debug('> {0} : get task pages ({1} tasks per page) randomly (0 < page < {2}) for {3} seconds'
                      .format(i, self.page_size, nb_pages, self.timeout, i))

        while time.time() < walltime:
            start = random.randint(0, min(nb_tasks, max_page))
            nb_accesses += 1
            try:
                tasks = t_local.tasks.list(sort='num' if sorted else None,
                                           dir='DESC',
                                           start=start,
                                           limit=self.page_size,
                                           filters=[],
                                           count=count,
                                           raw=raw,
                                           cast=False)
            except Exception as e:
                self.logger.warning('Cannot find task page {} / {}  : {}'.format(start, nb_tasks, e))
                nb_errors += 1

        self.logger.debug('+ {4} : {0} monitoring accesses in {2} seconds, {1} errors ({3} accesses per second)'
              .format(nb_accesses, nb_errors, self.timeout, nb_accesses / self.timeout, i))

        q.put([nb_accesses, nb_errors])


    def access_task_page_mongo_count(self, q, i):
        self.access_task_page_mongo(q, i, count=True)

    def access_task_page_mongo(self, q, i, max_page=100, sorted=True, count=False):

        nb_accesses = 0
        nb_errors = 0
        mask = {'masks.archived': None, 'masks.draft': None}

        client = MongoClient(self.mongo_url)
        db = client.Torii
        
        nb_tasks = db.tasks.find(mask).count()
        nb_pages = nb_tasks // self.page_size

        walltime = time.time() + self.timeout

        self.logger.debug('> {0} : get task pages ({1} tasks per page) randomly (0 < page < {2}) for {3} seconds'
                      .format(i, self.page_size, nb_pages, self.timeout, i))

        while time.time() < walltime:
            start = random.randint(0, min(nb_tasks, max_page))
            nb_accesses += 1
            try:
                if count:
                    nb_tasks = db.tasks.find(mask).count()

                if sorted:
                    tasks = db.tasks.find(mask).sort([('num', -1)]).skip(start).limit(self.page_size)
                else:
                    tasks = db.tasks.find(mask).skip(start).limit(self.page_size)

                # make sure to consume the cursor tasks (effective mongo access)
                for task in tasks:
                    pass

            except Exception as e:
                self.logger.warning('Cannot find task page {} / {}  : {}'.format(start, nb_tasks, e))
                nb_errors += 1

        self.logger.debug('+ {4} : {0} monitoring accesses in {2} seconds, {1} errors ({3} accesses per second)'
              .format(nb_accesses, nb_errors, self.timeout, nb_accesses / self.timeout, i))

        q.put([nb_accesses, nb_errors])


    def access_single_tasks(self, q, i):
        nb_accesses = 0
        nb_errors = 0

        t_local = Torii(base_url=self.torii_url, username='hpcgadmin', password='hpcgadmin')

        task = t_local.tasks.list(dir='ASC', sort='num', limit=1)[0]
        min_num = task.num
        task = t_local.tasks.list(dir='DESC', sort='num', limit=1)[0]
        max_num = task.num

        walltime = time.time() + self.timeout

        self.logger.debug('> {3} : get one task at a time randomly ({0} < num < {1}) for {2} seconds'
                      .format(min_num, max_num, self.timeout, i))

        while time.time() < walltime:
            num = random.randint(min_num, max_num)
            nb_accesses += 1
            try:
                no_archived_filter = {'property':'tags','value':'Test','filterType':'string','operator':'like'}
                task = t_local.tasks.get(num=num, take_first=True, filters=[no_archived_filter])
            except Exception as e:
                self.logger.warning('Cannot find task with num={0} : {1}'.format(num, e))
                nb_errors += 1

        self.logger.debug('+ {4} : {0} monitoring accesses in {2} seconds, {1} errors ({3} accesses per second)'
              .format(nb_accesses, nb_errors, self.timeout, nb_accesses // self.timeout, i))

        q.put([nb_accesses, nb_errors])


    def access_single_tasks_mongo(self, q, i):
        nb_accesses = 0
        nb_errors = 0

        client = MongoClient(self.mongo_url)
        db = client.Torii
        task = next(db.tasks.find({}).sort([('num', 1)]).limit(1))
        min_num = task['num']
        task = next(db.tasks.find({}).sort([('num', -1)]).limit(1))
        max_num = task['num']

        walltime = time.time() + self.timeout

        self.logger.debug('> {3} : get one task at a time randomly ({0} < num < {1}) for {2} seconds'
                         .format(min_num, max_num, self.timeout, i))

        while time.time() < walltime:
            num = random.randint(min_num, max_num)
            nb_accesses += 1
            try:
                task = next(db.tasks.find({'num': num}).limit(1))
            except Exception as e:
                self.logger.warning('Cannot find task with num={0} : {1}'.format(num, e))
                nb_errors += 1

        self.logger.debug('+ {4} : {0} monitoring accesses in {2} seconds, {1} errors ({3} accesses per second)'
              .format(nb_accesses, nb_errors, self.timeout, nb_accesses / self.timeout, i))

        q.put([nb_accesses, nb_errors])


    def run(self, nb_procs, method, use_case=None, ref_data=None):

        # run accepts arrays of nb_procs
        if isinstance(nb_procs, list):
            # run on each nb_proc value independently
            for value in nb_procs:
                self.run(value, method, use_case, ref_data)
            return

        # run accepts arrays of method
        if isinstance(method, list):
            # run on each method value independently
            for value in method:
                self.run(nb_procs, value, use_case, ref_data)
            return

        # do the actual run
        done = False
        try_num = 0
        while not done and try_num < self.nb_retry:
            try :
                try_num += 1
                self._run(nb_procs, method, use_case, ref_data)
                done = True
            except Exception as e:
                self.logger.error("Use case {} failed ({}th try) : {}"
                                  .format(use_case if use_case else '<anonymous>', try_num, e))

        if not done:
            self.logger.error("Use case {} failed definitively after {} tries"
                              .format(use_case if use_case else '<anonymous>', try_num))
            exit(-1)



    def _run(self, nb_procs, method, use_case=None, ref_data=None):

        self.nb_run += 1
        if not use_case:
            use_case = 'benchmark ' + str(self.nb_run)

        self.logger.info('==============================================================================')

        self.logger.info(' Benchmark \'{}\' (method={}, nb_procs={})'.format(use_case, method, nb_procs))
        self.logger.info('')

        self.logger.info('Generate random access to the tasks for {0} seconds with {1} processes'
              .format(self.timeout, nb_procs))

        procs = [None] * nb_procs
        q = Queue()

        for i in range(len(procs)):
            procs[i] = Process(target=getattr(self, method),
                               args=(q, i))
            procs[i].start()

        total_accesses = 0
        total_errors = 0
        for proc in procs:
            result = q.get(timeout=self.timeout * 2)
            total_accesses += result[0]
            total_errors += result[1]

        for proc in procs:
            proc.join()

        self.logger.info('= {0} monitoring accesses in {2} seconds, {1} errors ({3} accesses per second)'
              .format(total_accesses, total_errors, self.timeout, total_accesses / self.timeout))
        self.logger.info('==============================================================================')
        self.logger.info('')

        # ' num | use case | method | nb_procs | ref 1 | ref 2 | nb reqs | nb errors | time(s) | req/s  | date )\n')
        self.csv_file.write(' {:10d} | {:50s} | {:30s} | {:10d} |'.format(
            self.nb_run,
            use_case,
            method,
            nb_procs
        ))

        # write all the reference values
        if isinstance(ref_data, dict):
            for key in self.ref_data_keys:
                if key in ref_data:
                    self.csv_file.write(' {:20s} |'.format(str(ref_data[key])))

        self.csv_file.write(' {:10d} | {:10d} | {:10.1f} | {:10.1f} | {:10d} |  {:30s} |\n'.format(
            total_accesses,
            total_errors,
            self.timeout,
            total_accesses / self.timeout,
            int(time.time()),
            datetime.datetime.today().strftime('%d/%m/%Y %H:%M:%S')
        ))
        self.csv_file.flush()

if __name__ == '__main__':

    bench = MonitoringBenchmark(mongo_url='mongodb://10.140.9.172:27017',
                                torii_url='https://10.140.9.172:8443/torii',
                                timeout=60,
                                page_size=20)


    nb_procs = [1, 2, 3, 4, 5, 10, 15, 20, 30, 40]
    method = ['access_single_tasks', 'access_task_page', 'access_single_tasks_mongo']

    bench.run(nb_procs, method)

