
import threading
from queue import Queue
from multiprocessing.pool import ApplyResult

import tabular_logger as tlogger

class AsyncWorker(object):
    @property
    def concurrent_tasks(self):
        raise NotImplementedError()

    def run_async(self, task_id, task, callback):
        raise NotImplementedError()


class WorkerHub(object):
    def __init__(self, workers, input_queue, done_queue):
        self.done_buffer = Queue()
        self.workers = workers
        self.available_workers = Queue()
        self.done_queue = done_queue
        self._cache = {}
        self.input_queue = input_queue

        for w in workers:
            for t in w.concurrent_tasks:
                self.available_workers.put((w, t))

        self.__initialize_handlers()

    def __initialize_handlers(self):
        self._input_handler = threading.Thread(
            target=WorkerHub._handle_input,
            args=(self,)
            )
        self._input_handler._state = 0
        tlogger.info('WorkerHub: _input_handler initialized')

        self._output_handler = threading.Thread(
            target=WorkerHub._handle_output,
            args=(self,)
            )
        self._output_handler._state = 0
        tlogger.info('WorkerHub: _output_handler initialized')

    def worker_callback(self, worker, subworker, result):
        worker_task = (worker, subworker)
        if worker_task in self._cache:
            task_id = self._cache[worker_task]
            del self._cache[worker_task]
            self.done_buffer.put((task_id, result))
        else:
            tlogger.warn('WorkerHub: Worker task not found in cache', worker_task)
            tlogger.warn('WorkerHub: Subworker', subworker)
            tlogger.warn('WorkerHub: Unable to process result', result)

        # Return worker back
        self.available_workers.put(worker_task)

    @staticmethod
    def _handle_input(self):
        try:
            while True:
                worker_task = self.available_workers.get()
                if worker_task is None:
                    tlogger.info('WorkerHub._handle_input NO MORE WORKERS AWAILABLE')
                    break
                worker, subworker = worker_task

                task = self.input_queue.get()
                if task is None:
                    tlogger.info('WorkerHub._handle_input NO MORE INPUTS AWAILABLE')
                    break
                task_id, task = task
                self._cache[worker_task] = task_id
                # tlogger.info('WorkerHub: put task id: %s in cache keyed by worker task: %s' % (task_id, worker_task))

                worker.run_async(subworker, task, callback=self.worker_callback)
        except:
            tlogger.exception('WorkerHub._handle_input exception thrown')
            raise

    @staticmethod
    def _handle_output(self):
        try:
            while True:
                result = self.done_buffer.get()
                if result is None:
                    tlogger.info('WorkerHub._handle_output done')
                    break
                self.done_queue.put(result)
        except:
            tlogger.exception('WorkerHub._handle_output exception thrown')
            raise

    def initialize(self):
        self._input_handler.start()
        self._output_handler.start()

    def close(self):
        self.available_workers.put(None)
        self.input_queue.put(None)
        self.done_buffer.put(None)

class AsyncTaskHub(object):
    def __init__(self, input_queue=None, results_queue=None):
        if input_queue is None:
            input_queue = Queue(64)
        self.input_queue = input_queue
        self._cache = {}
        self.results_queue = None
        if results_queue is not None:
            self.results_queue = results_queue

            self._output_handler = threading.Thread(
                target=AsyncTaskHub._handle_output,
                args=(self,)
                )
            self._output_handler.daemon = True
            self._output_handler._state = 0
            self._output_handler.start()

    @staticmethod
    def _handle_output(self):
        try:
            while True:
                result = self.results_queue.get()
                if result is None:
                    tlogger.info('AsyncTaskHub._handle_output done')
                    break
                self.put(result)
        except:
            tlogger.exception('AsyncTaskHub._handle_output exception thrown')
            raise

    def run_async(self, task, callback=None, error_callback=None):
        result = ApplyResult(self._cache, callback, error_callback)
        self.input_queue.put((result._job, task))
        return result

    def put(self, result):
        job, result=result
        self._cache[job]._set(0, (True, result))


[((430, 437), 'queue.Queue', 'Queue', ([], {}), '()\n', (435, 437), False, 'from queue import Queue\n'), ((502, 509), 'queue.Queue', 'Queue', ([], {}), '()\n', (507, 509), False, 'from queue import Queue\n'), ((836, 898), 'threading.Thread', 'threading.Thread', ([], {'target': 'WorkerHub._handle_input', 'args': '(self,)'}), '(target=WorkerHub._handle_input, args=(self,))\n', (852, 898), False, 'import threading\n'), ((984, 1037), 'tabular_logger.info', 'tlogger.info', (['"""WorkerHub: _input_handler initialized"""'], {}), "('WorkerHub: _input_handler initialized')\n", (996, 1037), True, 'import tabular_logger as tlogger\n'), ((1070, 1133), 'threading.Thread', 'threading.Thread', ([], {'target': 'WorkerHub._handle_output', 'args': '(self,)'}), '(target=WorkerHub._handle_output, args=(self,))\n', (1086, 1133), False, 'import threading\n'), ((1220, 1274), 'tabular_logger.info', 'tlogger.info', (['"""WorkerHub: _output_handler initialized"""'], {}), "('WorkerHub: _output_handler initialized')\n", (1232, 1274), True, 'import tabular_logger as tlogger\n'), ((4565, 4615), 'multiprocessing.pool.ApplyResult', 'ApplyResult', (['self._cache', 'callback', 'error_callback'], {}), '(self._cache, callback, error_callback)\n', (4576, 4615), False, 'from multiprocessing.pool import ApplyResult\n'), ((1581, 1651), 'tabular_logger.warn', 'tlogger.warn', (['"""WorkerHub: Worker task not found in cache"""', 'worker_task'], {}), "('WorkerHub: Worker task not found in cache', worker_task)\n", (1593, 1651), True, 'import tabular_logger as tlogger\n'), ((1664, 1711), 'tabular_logger.warn', 'tlogger.warn', (['"""WorkerHub: Subworker"""', 'subworker'], {}), "('WorkerHub: Subworker', subworker)\n", (1676, 1711), True, 'import tabular_logger as tlogger\n'), ((1724, 1783), 'tabular_logger.warn', 'tlogger.warn', (['"""WorkerHub: Unable to process result"""', 'result'], {}), "('WorkerHub: Unable to process result', result)\n", (1736, 1783), True, 'import tabular_logger as tlogger\n'), ((3591, 3600), 'queue.Queue', 'Queue', (['(64)'], {}), '(64)\n', (3596, 3600), False, 'from queue import Queue\n'), ((3820, 3886), 'threading.Thread', 'threading.Thread', ([], {'target': 'AsyncTaskHub._handle_output', 'args': '(self,)'}), '(target=AsyncTaskHub._handle_output, args=(self,))\n', (3836, 3886), False, 'import threading\n'), ((2715, 2776), 'tabular_logger.exception', 'tlogger.exception', (['"""WorkerHub._handle_input exception thrown"""'], {}), "('WorkerHub._handle_input exception thrown')\n", (2732, 2776), True, 'import tabular_logger as tlogger\n'), ((3128, 3190), 'tabular_logger.exception', 'tlogger.exception', (['"""WorkerHub._handle_output exception thrown"""'], {}), "('WorkerHub._handle_output exception thrown')\n", (3145, 3190), True, 'import tabular_logger as tlogger\n'), ((4396, 4461), 'tabular_logger.exception', 'tlogger.exception', (['"""AsyncTaskHub._handle_output exception thrown"""'], {}), "('AsyncTaskHub._handle_output exception thrown')\n", (4413, 4461), True, 'import tabular_logger as tlogger\n'), ((2066, 2131), 'tabular_logger.info', 'tlogger.info', (['"""WorkerHub._handle_input NO MORE WORKERS AWAILABLE"""'], {}), "('WorkerHub._handle_input NO MORE WORKERS AWAILABLE')\n", (2078, 2131), True, 'import tabular_logger as tlogger\n'), ((2306, 2370), 'tabular_logger.info', 'tlogger.info', (['"""WorkerHub._handle_input NO MORE INPUTS AWAILABLE"""'], {}), "('WorkerHub._handle_input NO MORE INPUTS AWAILABLE')\n", (2318, 2370), True, 'import tabular_logger as tlogger\n'), ((2984, 3029), 'tabular_logger.info', 'tlogger.info', (['"""WorkerHub._handle_output done"""'], {}), "('WorkerHub._handle_output done')\n", (2996, 3029), True, 'import tabular_logger as tlogger\n'), ((4260, 4308), 'tabular_logger.info', 'tlogger.info', (['"""AsyncTaskHub._handle_output done"""'], {}), "('AsyncTaskHub._handle_output done')\n", (4272, 4308), True, 'import tabular_logger as tlogger\n')]