import os
import time
import tempfile
from threading import Thread

import datetime

from torii import Torii


class LowlaBenchmark(object):


    def __init__(self,
                 mongo_url='mongodb://localhost:27017',
                 torii_url='http://localhost:8443/torii',
                 timeout=300,
                 nb_retry=3,
                 output_file=None,
                 nb_threads=10,
                 nb_users=10,
                 nb_lowla_tasks=1000,
                 ref_data_keys=[]):

        self.mongo_url = mongo_url
        self.torii_url = torii_url
        self.timeout = timeout
        self.nb_run = 0
        self.nb_retry = nb_retry
        self.ref_data_keys = ref_data_keys
        self.nb_threads = nb_threads
        self.nb_users= nb_users
        self.nb_lowla_tasks= nb_lowla_tasks

        if output_file:
            self.output_file = output_file
        else:
            self.output_file = 'torii_benchmark_{}_submit.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.t_users = []

        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 Submit Benchmark \n'
                            '## {0}\n'
                            '##\n\n\n'.format(datetime.datetime.today().strftime('%Y-%m-%d %H:%M:%S')))

        # ' num | use case | nb_threads | ref 1 | ref 2 | nb_tasks | create time | sched time | submit time | total time | tasks/s  | date |)
        self.csv_file.write(' num        |'
                            ' use case                                           |'
                            ' nb users   |'
                            ' nb threads |')

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

        self.csv_file.write(' nb tasks        |'
                            ' lowla time      |'
                            ' task/sec        |'
                            ' timestamp       |'
                            ' date                           |\n')

        self.csv_file.flush()

        # finish the preparation of the bench
        self.prepare_bench()


    def get_t(self, i):
        return self.t_users[i % len(self.t_users)]

    def create_tasks(self, tasks, index):
        t = self.get_t(index)
        begin = time.time()
    
        # create all the tasks
        task = t.tasks.new(name='lowla engine {0}'.format(index),
                               application=self.application)
        task.set_input('sleep', 0)
        task.set_cluster(self.cluster)
        task.set_rundir(tempfile.mkdtemp())
        task.submit()
        task.wait(status='RUNNING')

        tasks[index] = task
    
        end = time.time()

        self.logger.info('\t+ engine task running in {0} seconds'.format(end - begin))
    
    
    def submit_lowla_tasks(self, nb_lowla_tasks, tasks, index):
        t = self.get_t(index)
        task = tasks[index]
        begin = time.time()
        for i in range(nb_lowla_tasks):
            try:
                task.lowla()
            except Exception as e:
                self.logger.error('! {} : cannot submit lowla task on {} as {} : {}'
                                  .format(index, task.num, t.session.user.name, e))

        end = time.time()

        self.logger.info('\t+ {0} lowla tasks are submitted in {1} seconds'.format(nb_lowla_tasks, end - begin))
    

    def delete_tasks(self, tasks, index):
        t = self.get_t(index)

        begin = time.time()

        for task in tasks:
            if task.status == 'REGISTERED':
                t.tasks.delete(task)
    
        end = time.time()

        self.logger.info('\t+ {0} tasks are delete in {1} seconds'.format(len(tasks), end - begin))

    def run(self, nb_lowla_tasks=None, nb_threads=None, ref_data=None, use_case=None):
        
        # run accepts arrays of nb_tasks
        if isinstance(nb_lowla_tasks, list):
            # run on each nb_proc value independently
            for value in nb_lowla_tasks:
                self.run(nb_lowla_tasks=value, nb_threads=nb_threads)
            return

        # run accepts arrays of nb_threads
        if isinstance(nb_threads, list):
            # run on each nb_proc value independently
            for value in nb_threads:
                self.run(nb_lowla_tasks=nb_lowla_tasks, nb_threads=value)
            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_lowla_tasks=nb_lowla_tasks, nb_threads=nb_threads, ref_data=ref_data, use_case=None)
                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_lowla_tasks=None, nb_threads=None, ref_data=None, use_case=None):

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

        if not nb_lowla_tasks:
            nb_lowla_tasks = self.nb_lowla_tasks
        if not nb_threads:
            nb_threads = self.nb_threads

        threads = [None] * nb_threads
        tasks = [None] * nb_threads

        # Initiate the sessions for all the users
        self.t_users = []
        for i in range(self.nb_threads):
            username = 'hpcg{0}'.format(i % self.nb_users + 1)
            self.t_users.append(Torii(base_url=self.torii_url, username=username, password=username))

        # if nb_threads > len(self.t_users):
        #     for i in range(len(self.t_users), self.nb_threads):
        #         username = 'testhpcg{0}'.format(i % self.nb_users + 1)
        #         self.t_users.append(Torii(base_url=self.torii_url, username=username, password=username))


        # each user creates his share of the tasks
        self.logger.info('Create all the engine tasks')
        begin = time.time()
        for i in range(nb_threads):
            threads[i] = Thread(target=self.create_tasks,
                                args=(tasks, i))
            threads[i].start()

        for thread in threads:
            thread.join()

        end = time.time()
        create_time = end - begin
        self.logger.info('= {0} engine tasks are created in {1} seconds ({2} tasks per second)'
                         .format(nb_threads, end - begin, nb_threads / (end - begin)))


        # each user creates his share of the tasks
        self.logger.info('')
        begin = time.time()
        for i in range(nb_threads):
            threads[i] = Thread(target=self.submit_lowla_tasks,
                                args=(nb_lowla_tasks//nb_threads, tasks, i))
            threads[i].start()

        for thread in threads:
            thread.join()

        end = time.time()
        self.logger.info('= {0} lowla tasks are executed in {1} seconds ({2} tasks per second)'
                         .format(nb_lowla_tasks, end - begin, nb_lowla_tasks / (end - begin)))
        self.logger.info('')
        create_time = end - begin

        # each user waits his own tasks
        begin = time.time()
        for i in range(nb_threads):
            threads[i] = Thread(target=self.delete_tasks, args=(tasks, i))
            threads[i].start()

        for thread in threads:
            thread.join()

        end = time.time()
        self.logger.info('= {0} engine tasks are deleted in {1} seconds ({2} tasks per second)'
              .format(nb_threads, end - begin, nb_threads / (end - begin)))



    def prepare_bench(self):
        self.logger.info('Get the application and cluster to submit too and make sure testusers can use them')

        application = self.t.applications.get(name='ToriiTest LowLa')
        cluster = self.t.clusters.get(type='hpc', take_first=True)
        server = self.t.servers.get(cluster.headnode.id)
        assert application
        assert cluster
        assert server
    
        user = self.t.users.get(name='hpcg1')
        assert user and user.teams
    
        cluster.teams = user.teams
        server.teams = user.teams
        application.teams = user.teams

        self.cluster = cluster.update()
        self.application = application.update()
        self.server = server.update()

        self.t_users = []


if __name__ == '__main__':

    bench = LowlaBenchmark(mongo_url='mongodb://10.140.9.23:26017',
                            torii_url='https://10.140.9.23:6443/torii',
                            nb_threads=10,
                            nb_users=2,
                            nb_lowla_tasks=100,
                            timeout=300)

    panel = [1, 2, 4, 8, 16, 32]
    panel = [32]
    for nb_threads in panel:
        bench.run(nb_threads=nb_threads, nb_lowla_tasks=nb_threads * 1000)


