from .rate import RateTracker
from .worker import ThreadWorker
from .task import Task, RatedTask
from collections import deque
from threading import Lock
import functools
import logging

log = logging.getLogger("ThreadWorkerQueue")


class ThreadWorkerQueue:
    instance = None
    RATE_GROUP = {}

    # TODO! should we increase io_worker_limit when we've doubled vCPUs in fargate?
    def __init__(self, thread_count=None, io_worker_limit=10):  # threadCount = None = detectedCPU's*cores +1
        self._workPool = deque()
        self._workers = []
        self._io_workPool = deque()
        self._io_workers = []
        self.io_worker_limit = io_worker_limit
        self.lock = Lock()
        if thread_count is None:
            # figure out real CPU count
            try:
                from multiprocessing import cpu_count
                thread_count = cpu_count() + 1  # Comes out, +1 is just slowing things down..
            except Exception as e:
                log.exception(e)
                thread_count = 2

        for w in range(thread_count):
            self._workers.append(ThreadWorker(self._workPool, f"CPU_Worker-{w}"))

    def __del__(self):
        self.stop()

    def stop(self):
        wpl = len(self._workPool) + len(self._io_workPool)
        if wpl > 0:
            log.warning(f"ThreadWorkerQueue.stop: There were some tasks left in queue: {wpl}")
        for worker in self._workers:
            if worker.is_busy():
                log.warning(f"ThreadWorkerQueue.stop: a worker was busy")
            worker.stop()
        del (self._workers[:])

        for worker in self._io_workers:
            if worker.is_busy():
                log.warning(f"ThreadWorkerQueue.stop: a worker was busy")
            worker.stop()
        del (self._io_workers[:])

    @staticmethod
    def instantiate(thread_count=2, io_worker_limit=10):
        if ThreadWorkerQueue.instance is None:
            ThreadWorkerQueue.instance = ThreadWorkerQueue(thread_count=thread_count, io_worker_limit=io_worker_limit)  #
        return ThreadWorkerQueue.instance

    @staticmethod
    def terminate():
        # Because thread workers hold the reference to the workerqueue, they have to be released first.
        # Otherwise the instance will not be released.
        if ThreadWorkerQueue.instance:
            ThreadWorkerQueue.instance.stop()
            ThreadWorkerQueue.instance = None

    def idle_workers(self):
        if self._workers:
            for worker in self._workers:
                if worker.is_idle():
                    yield worker

    def more_work(self, work):
        self._workPool.append(work)
        # log("DDThreadWorkerQueue: gotWork[{0}]".format(id(work)))
        # if there are not busy workers, tell one of them about it..
        for worker in self.idle_workers():
            worker.signal_work_arrived()
            break

    def io_work(self, work):
        """ Just handles io tasks, that can be plentiful.
        Reuses already created workers or just creates new, if all are busy.
        """
        self._io_workPool.append(work)
        for worker in self._io_workers:
            if worker.is_idle():
                worker.signal_work_arrived()
                return
        if len(self._io_workers) < self.io_worker_limit:
            with self.lock:  # make sure thread creation proceeds correctly.
                wl = len(self._io_workers)
                if wl < self.io_worker_limit:
                    new_worker = ThreadWorker(self._io_workPool, f"IO_Worker-{wl}")
                    self._io_workers.append(new_worker)
                    new_worker.signal_work_arrived()

    def _more_work(self, work):
        """ Internal version used by the threads. No need to notify other threads from workerthread """
        self._workPool.append(work)

    # decorator for worker methods
    @staticmethod
    def task(async_function):
        @functools.wraps(async_function)
        def start_working(*args, **kwargs):
            work = Task(async_function, args, kwargs)
            ThreadWorkerQueue.instantiate().more_work(work)
            return work

        return start_working

    @staticmethod
    def io_task(async_function):
        """Executes function asynchronously. Returns task_handle you can use to wait for result.
        Calls function with original arguments, adding task_handle parameter

        def task(*args, task_handle=Task, **kwargs)"""
        @functools.wraps(async_function)
        def start_working(*args, **kwargs):
            work = Task(async_function, args, kwargs)
            ThreadWorkerQueue.instantiate().io_work(work)
            return work

        return start_working

    @staticmethod
    def io_task_limited(rate: float, group: str = None):
        """Executes function asynchronously. Returns task_handle you can use to wait for result.
        Calls function with original arguments, adding task_handle parameter
        Rate limits the calling of the function, no matter how fast the results are coming in.

        :argument rate max number of calls per second
        :argument group calls in same group are limited with same rate. Rate of first group mention is used """

        def io_task(async_function):
            if group is not None:
                rate_tracker = ThreadWorkerQueue.RATE_GROUP.setdefault(group, RateTracker(rate))
            else:
                rate_tracker = RateTracker(rate)

            @functools.wraps(async_function)
            def start_working(*args, **kwargs):
                work = Task(async_function, args, kwargs)
                rate_tracker.create_wait()
                ThreadWorkerQueue.instantiate().io_work(work)
                return work

            return start_working

        return io_task


# To avoid using lock
# ThreadWorkerQueue.instantiate()