
from __future__ import absolute_import
from __future__ import print_function

import itertools
import logging
import threading

from six.moves import queue
from time import sleep

from vdsm.common import concurrent


class ThreadPool:

    log = logging.getLogger('storage.ThreadPool')

    def __init__(self, name, numThreads, waitTimeout=3, maxTasks=100):

        self.log.debug("Enter - name: %s, numThreads: %s, waitTimeout: %s, "
                       "maxTasks: %s",
                       name, numThreads, waitTimeout, maxTasks)
        self._name = name
        self._count = itertools.count()
        self.__threads = []
        self._taskThread = {}
        self.__resizeLock = threading.Condition(threading.Lock())
        self.__runningTasksLock = threading.Condition(threading.Lock())
        self.__tasks = queue.Queue(maxTasks)
        self.__isJoining = False
        self.__runningTasks = 0
        self.__waitTimeout = waitTimeout
        self.setThreadCount(numThreads)

    def setRunningTask(self, addTask):

        self.__runningTasksLock.acquire()
        try:
            if addTask:
                self.__runningTasks += 1
            else:
                self.__runningTasks -= 1
            self.log.debug("Number of running tasks: %s", self.__runningTasks)
        finally:
            self.__runningTasksLock.release()

    def getRunningTasks(self):
        return self.__runningTasks

    def setThreadCount(self, newNumThreads):

        if self.__isJoining:
            return False

        self.__resizeLock.acquire()
        try:
            self.__setThreadCountNolock(newNumThreads)
        finally:
            self.__resizeLock.release()
        return True

    def __setThreadCountNolock(self, newNumThreads):

        while newNumThreads > len(self.__threads):
            name = "%s/%d" % (self._name, next(self._count))
            newThread = WorkerThread(self, name)
            self.__threads.append(newThread)
            newThread.start()
        while newNumThreads < len(self.__threads):
            self.__threads[0].goAway()
            del self.__threads[0]

    def getThreadCount(self):

        self.__resizeLock.acquire()
        try:
            return len(self.__threads)
        finally:
            self.__resizeLock.release()

    def queueTask(self, id, task, args=None, taskCallback=None):

        if self.__isJoining:
            return False
        if not callable(task):
            return False

        self.__tasks.put((id, task, args, taskCallback))

        return True

    def getNextTask(self):

        id = None
        cmd = None
        args = None
        callback = None

        try:
            id, cmd, args, callback = self.__tasks.get(True,
                                                       self.__waitTimeout)
        except queue.Empty:
            pass

        return id, cmd, args, callback

    def stopThread(self):
        return self.__tasks.put((None, None, None, None))

    def joinAll(self, waitForTasks=True, waitForThreads=True):

        self.__isJoining = True

        if waitForTasks:
            while not self.__tasks.empty():
                sleep(0.1)

        self.__resizeLock.acquire()
        try:
            if waitForThreads:
                for t in self.__threads:
                    t.goAway()
                for t in self.__threads:
                    t.join()
                    del t
            self.__setThreadCountNolock(0)
            self.__isJoining = True

            self.__isJoining = False
        finally:
            self.__resizeLock.release()


class WorkerThread(object):

    log = logging.getLogger('storage.ThreadPool.WorkerThread')

    def __init__(self, pool, name):

        self._thread = concurrent.thread(self.run, name=name)
        self.__pool = pool
        self.__isDying = False

    def start(self):
        self._thread.start()

    def join(self):
        self._thread.join()

    def _processNextTask(self):
        id, cmd, args, callback = self.__pool.getNextTask()
        try:
                pass
            elif self.__isDying:
                self.__pool.__tasks.put((id, cmd, args, callback))
            elif callback is None:
                self.__pool.setRunningTask(True)
                self.log.info("START task %s (cmd=%r, args=%r)",
                              id, cmd, args)
                cmd(args)
                self.log.info("FINISH task %s", id)
                self.__pool.setRunningTask(False)
            else:
                self.__pool.setRunningTask(True)
                self.log.info("START task %s (callback=%r, cmd=%r, args=%r)",
                              id, callback, cmd, args)
                callback(cmd(args))
                self.log.info("FINISH task %s", id)
                self.__pool.setRunningTask(False)
        except Exception:
            self.log.exception("FINISH task %s failed (callback=%r, "
                               "cmd=%r, args=%r)",
                               id, callback, cmd, args)

    def run(self):

        while not self.__isDying:
            self._processNextTask()

    def goAway(self):

        self.__isDying = True


if __name__ == "__main__":

    from random import randrange


    def sortTask(data):
        print("SortTask starting for ", data)
        numbers = list(range(data[0], data[1]))
        for a in numbers:
            rnd = randrange(0, len(numbers) - 1)
            a, numbers[rnd] = numbers[rnd], a
        print("SortTask sorting for ", data)
        numbers.sort()
        print("SortTask done for ", data)
        return "Sorter ", data


    def waitTask(data):
        print("WaitTask starting for ", data)
        print("WaitTask sleeping for %d seconds" % data)
        sleep(data)
        return "Waiter ", data


    def taskCallback(data):
        print("Callback called for", data)


    pool = ThreadPool("work", 100)

    print("Running tasks: ", pool.getRunningTasks(), "\n")
    pool.queueTask(sortTask, (1000, 100000), taskCallback)
    print("Running tasks: ", pool.getRunningTasks(), "\n")
    pool.queueTask(waitTask, 5, taskCallback)
    pool.queueTask(sortTask, (200, 200000), taskCallback)
    pool.queueTask(waitTask, 2, taskCallback)
    print("Running tasks: ", pool.getRunningTasks(), "\n")
    pool.queueTask(sortTask, (3, 30000), taskCallback)
    pool.queueTask(waitTask, 7, taskCallback)

    pool.joinAll()
    print("ThreadPool sample done.\n")
