import zmq
import time
from typing import Any, List, ClassVar
from threading import Thread

from ._WorkerThread import WorkerThread
from ._ThreadPoolParameters import ThreadPoolParameters


class ZeroMqPushPullThreadPool(object):
    def __init__(self, context, pull_address, parameters, worker_address, thread_class, *args, **kwargs):
        # type: (Any, str, ThreadPoolParameters, str, ClassVar, Any, Any) -> None
        self.context = context
        self.pull_address = pull_address
        self.parameters = parameters
        self.worker_address = worker_address
        self.thread_class = thread_class
        self.args = args
        self.kwargs = kwargs

        self.threads = []  # type: List[WorkerThread]
        self.pull_socket = context.socket(zmq.PULL)
        self.pull_socket.bind(self.pull_address)

        self.push_socket = context.socket(zmq.PUSH)
        self.push_socket.bind(self.worker_address)

        for i in range(0, self.parameters.initial_count, 1):
            self.create_worker_thread()

        self.monitor_thread = Thread(target=self.monitor)
        self.monitor_thread.daemon = True
        self.monitor_thread.start()

        self.running = True
        self.pull_push_thread = Thread(target=self.pull_push)
        self.pull_push_thread.start()

    def get_current_thread_count(self):
        return len(self.threads)

    def dispose(self):
        self.running = False
        # self.pull_push_thread.join()

        self.pull_socket.close()

        for thread in self.threads:
            thread.stop()

        for thread in self.threads:
            thread.join()

        self.threads.clear()
        self.monitor_thread.join()
        self.push_socket.close()

    def pull_push(self):
        while self.running:
            ready = self.pull_socket.poll(1000, zmq.POLLIN)
            if ready != zmq.POLLIN:
                continue

            data = self.pull_socket.recv()
            self.push_socket.send(data)

    def create_worker_thread(self):
        args = [self.context, self.worker_address, self.parameters.timeout]
        args.extend(self.args)
        thread = self.thread_class(*args, **self.kwargs)
        thread.start()
        self.threads.append(thread)

    def monitor(self):
        while len(self.threads) > 0:
            time.sleep(self.parameters.monitor_timeout)
            busy_threads_count = 0
            idle_threads = []
            for thread in self.threads:
                if thread.is_working:
                    busy_threads_count += 1
                elif thread.is_idle:
                    idle_threads.append(thread)

            if busy_threads_count == len(self.threads) and len(self.threads) < self.parameters.max_count:
                increment = self.parameters.increment \
                    if len(self.threads) + self.parameters.increment <= self.parameters.max_count \
                    else self.parameters.max_count - len(self.threads)

                for i in range(0, increment, 1):
                    self.create_worker_thread()

            if len(self.threads) > self.parameters.initial_count and len(idle_threads) > 0:
                for thread in idle_threads:
                    thread.stop()
                    self.threads.remove(thread)

                    if len(self.threads) == self.parameters.initial_count:
                        break
