import time
from typing import Dict

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
from itertools import chain
from time import sleep
log = logging.getLogger().getChild("ThreadWorkerQueue")


class ThreadWorkerQueue:
    instance = None
    RATE_GROUP: Dict[str, RateTracker] = {}

    # 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

    @staticmethod
    def wait_for_complete(timeout=None):
        if ThreadWorkerQueue.instance:
            everything_done = False
            start_time = time.time()
            while not everything_done and (timeout is None or (time.time() - start_time) < timeout):
                for worker in chain(ThreadWorkerQueue.instance._workers, ThreadWorkerQueue.instance._io_workers):
                    while worker.is_busy():
                        sleep(0.01)
                        continue
                if ThreadWorkerQueue.instance._workPool or ThreadWorkerQueue.instance._io_workPool:
                    sleep(0.01)
                    continue
                everything_done = True
            return everything_done

    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()