import copy
import time
import tempfile

from .monitoring_benchmark import MonitoringBenchmark
from .submit_benchmark import SubmitBenchmark

from multiprocessing import Queue, Process

from pymongo import MongoClient

from torii import Torii

TAGS = ['draft', 'hibou', 'caillou', 'chou', 'genou', 'bijou', 'joujou', 'pou']
ARCHIVED_TAG = 'archived'

class TaskGenerator(object):

    def __init__(self, nb_procs=20, nb_users=10):

        self.nb_procs = nb_procs
        self.nb_users = nb_users

        self.t = Torii(base_url=base_url, username='hpcgadmin', password='hpcgadmin')
        self.logger = self.t.logger
        self.logger.setLevel('INFO')

        self.logger.info('Gives rights to all the users on all the applications')
        user = self.t.users.get(name='testhpcg1')
        assert user and user.teams

        for cluster in self.t.clusters.list():
            cluster.teams = user.teams
            cluster.update()

        for server in self.t.servers.list():
            server.teams = user.teams
            server.update()

        for application in self.t.applications.list(name='Test'):
            application.teams = user.teams
            application.update()


    def create_tasks(self, t, nb_tasks, q, i):
        # Initiate the sessions for all the users
        # username = 'testhpcg{0}'.format(i % nb_users + 1)
        # t = Torii(base_url=base_url, username=username, password=username)


        cluster = t.clusters.get(type='hpc', take_first=True)
        applications = t.applications.list(name='Test App')

        begin = time.time()
        min_num = 0
        max_num = 0

        # create all the tasks
        for j in range(nb_tasks) :
            task = t.tasks.new(name='test {0} {1}'.format(i, j),
                               application=applications[j % len(applications)])
            task.tags.append(TAGS[j % len(TAGS)])
            task.set_cluster(cluster)
            task.set_rundir('/tmp/{}/{}'.format(i, begin))
            task.context.num_proc = i
            task.context.creation_rank = j
            task.update()

            # if j % 500 == 0 :
            #     t.logger.info('\t... {0} : {1} tasks are created so far'.format(i, j))

            if j == 0:
                min_num = task.num
                max_num = task.num
            elif task.num < min_num:
                min_num = task.num
            elif task.num > min_num:
                max_num = task.num

        end = time.time()

        t.logger.info('\t+ {0} : {1} tasks are created in {2} seconds'.format(i, nb_tasks, end - begin))

        q.put([min_num, max_num])



    def init_active_pool(self, nb_tasks):

        db = MongoClient(mongo_url).Torii

        current_size = db.tasks.count({'masks.' + ARCHIVED_TAG: {'$eq': True}})
        if current_size == nb_tasks:
            # nothing to do
            return
        elif current_size > nb_tasks:
            # In case we already pass the goal, we go back to the initial pool
            # clear the whole task collection and reset the task num increment
            db.tasks.delete_many({})
            db.inc.delete_many({})
            current_size = 0

        # do not recreate active tasks for nothing
        nb_tasks -= current_size

        procs = []
        q = Queue()

        self.logger.info('Create pool of {} "active" tasks (not archived)'.format(nb_tasks))
        # each user creates his share of the tasks
        begin = time.time()
        for i in range(nb_procs):

            proc_nb_tasks = nb_tasks // nb_procs if i > 0 else nb_tasks // nb_procs + nb_tasks % nb_procs
            if proc_nb_tasks > 0:
                # Initiate the sessions for all the users
                username = 'testhpcg{0}'.format(i % self.nb_users + 1)
                t_user = Torii(base_url=base_url, username=username, password=username)

                procs.append(Process(target=self.create_tasks,
                                     args=(t_user, proc_nb_tasks, q, i)))

        for proc in procs:
            proc.start()

        for i in range(len(procs)):
            result = q.get()
            min_num = result[0]
            max_num = result[1]

            if i == 0:
                min_num_global = min_num
                max_num_global = max_num
            elif min_num < min_num_global:
                min_num_global = min_num
            elif max_num > max_num_global:
                max_num_global = max_num


        for proc in procs:
            proc.join()

        end = time.time()
        global_time = end - begin
        self.logger.info('= {0} tasks are created in {1} seconds ({2} tasks per second)'
                            ', {3} < num < {4})'
              .format(nb_tasks, global_time, nb_tasks / global_time, min_num_global, max_num_global))
        self.logger.info('')


    def get_task_chunk(self, chunk_size):

        client = MongoClient(mongo_url)
        db = client.Torii

        tasks = []
        for task in db.tasks.find().limit(chunk_size):
            if ARCHIVED_TAG not in task['tags']:
                task['tags'].append(ARCHIVED_TAG)
                task['masks'] = {ARCHIVED_TAG: True}
            if task['num'] > 0:
                task['num'] = -task['num']
            del task['_id']

            tasks.append(task)

        self.logger.debug('Replication chunk prepared ({0} tasks)'.format(len(tasks)))
        return tasks


    def resize_collection(self, size, chunk_size):
        self.logger.debug('Resize task collection resized to {0} documents'.format(size))

        client = MongoClient(mongo_url)
        db = client.Torii

        current_size = db.tasks.count({})

        # In case we already pass the goal, we go back to the initial pool
        # remove all the retired
        if current_size > size:
            db.tasks.delete_many({'tags': {'$in': [ARCHIVED_TAG]}})
            current_size = client.Torii.tasks.count()

        while current_size < size:
            self.logger.debug('... {0} -> {1}'.format(current_size, size))
            chunk = self.get_task_chunk(chunk_size)
            db.tasks.insert_many(chunk)
            current_size = db.tasks.count({})

        self.logger.info('Task collection resized to {0} documents'.format(current_size))
        return current_size


def human_size(size):
    prefixes = ['K', 'M', 'G', 'T']

    prefix = ''
    i = 0
    hsize = float(size)
    while i < len(prefixes) and hsize >= 800:
        prefix = prefixes[i]
        hsize = hsize / float(1000)
        i += 1

    return '{0:.1f}{1}'.format(hsize, prefix).replace('.0', '')



if __name__ == '__main__':

    mongo_url = 'mongodb://10.140.9.172:27017'
    base_url = 'https://10.140.9.172:8443/torii'
    #mongo_url = 'mongodb://10.140.9.23:26017'
    #base_url = 'https://10.140.9.23:6443/torii'
    nb_tasks = 10000
    chunk_size = 2000
    max_size = 200000

    nb_procs = 10

    generator = TaskGenerator(nb_procs, nb_users=10)

    monitor_bench = MonitoringBenchmark(mongo_url=mongo_url,
                                        torii_url=base_url,
                                        timeout=20,
                                        page_size=20,
                                        ref_data_keys=['db_size'])

    submit_bench = SubmitBenchmark(mongo_url=mongo_url,
                                   torii_url=base_url,
                                   nb_threads=10,
                                   nb_users=10,
                                   nb_tasks=100,
                                   timeout=300,
                                   ref_data_keys=['db_size'])

    # methods = ['access_task_page', 'access_task_page_mongo', 'access_single_tasks', 'access_single_tasks_mongo']
    # methods = ['access_task_page_mongo_count', 'access_task_page_mongo']
    # methods = ['access_task_page', 'access_task_page_mongo_count', 'access_task_page_mongo']
    # methods = ['access_single_tasks', 'access_single_tasks_mongo']
    methods = ['access_task_page_count', 'access_task_page', 'access_task_page_raw', 'access_task_page_mongo']
    # methods = ['access_task_page_raw']

    # resize the pool of active tasks in the collection
    generator.init_active_pool(nb_tasks)

    current_size = nb_tasks
    while True:
        usecase = 'Bench accesses with {0} tasks in DB'.format(human_size(current_size))
        ref_data = {'db_size': current_size}

        # resize the whole collection (add only archived tasks)
        generator.resize_collection(current_size, chunk_size)

        # run the monitoring bench
        monitor_bench.run(nb_procs,
                          methods,
                          use_case=usecase,
                          ref_data=ref_data)

        # run the submission bench
        # submit_bench.run(use_case=usecase, ref_data=ref_data)

        if current_size > max_size:
            break;
        else:
            current_size *= 2
