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


[((19, 27, 19, 34), 'queue.Queue', 'Queue', ({}, {}), '()', False, 'from queue import Queue\n'), ((21, 33, 21, 40), 'queue.Queue', 'Queue', ({}, {}), '()', False, 'from queue import Queue\n'), ((33, 30, 36, 13), 'threading.Thread', 'threading.Thread', (), '', False, 'import threading\n'), ((38, 8, 38, 61), 'tabular_logger.info', 'tlogger.info', ({(38, 21, 38, 60): '"""WorkerHub: _input_handler initialized"""'}, {}), "('WorkerHub: _input_handler initialized')", True, 'import tabular_logger as tlogger\n'), ((40, 31, 43, 13), 'threading.Thread', 'threading.Thread', (), '', False, 'import threading\n'), ((45, 8, 45, 62), 'tabular_logger.info', 'tlogger.info', ({(45, 21, 45, 61): '"""WorkerHub: _output_handler initialized"""'}, {}), "('WorkerHub: _output_handler initialized')", True, 'import tabular_logger as tlogger\n'), ((138, 17, 138, 67), 'multiprocessing.pool.ApplyResult', 'ApplyResult', ({(138, 29, 138, 40): 'self._cache', (138, 42, 138, 50): 'callback', (138, 52, 138, 66): 'error_callback'}, {}), '(self._cache, callback, error_callback)', False, 'from multiprocessing.pool import ApplyResult\n'), ((54, 12, 54, 82), 'tabular_logger.warn', 'tlogger.warn', ({(54, 25, 54, 68): '"""WorkerHub: Worker task not found in cache"""', (54, 70, 54, 81): 'worker_task'}, {}), "('WorkerHub: Worker task not found in cache', worker_task)", True, 'import tabular_logger as tlogger\n'), ((55, 12, 55, 59), 'tabular_logger.warn', 'tlogger.warn', ({(55, 25, 55, 47): '"""WorkerHub: Subworker"""', (55, 49, 55, 58): 'subworker'}, {}), "('WorkerHub: Subworker', subworker)", True, 'import tabular_logger as tlogger\n'), ((56, 12, 56, 71), 'tabular_logger.warn', 'tlogger.warn', ({(56, 25, 56, 62): '"""WorkerHub: Unable to process result"""', (56, 64, 56, 70): 'result'}, {}), "('WorkerHub: Unable to process result', result)", True, 'import tabular_logger as tlogger\n'), ((109, 26, 109, 35), 'queue.Queue', 'Queue', ({(109, 32, 109, 34): '64'}, {}), '(64)', False, 'from queue import Queue\n'), ((116, 35, 119, 17), 'threading.Thread', 'threading.Thread', (), '', False, 'import threading\n'), ((81, 12, 81, 73), 'tabular_logger.exception', 'tlogger.exception', ({(81, 30, 81, 72): '"""WorkerHub._handle_input exception thrown"""'}, {}), "('WorkerHub._handle_input exception thrown')", True, 'import tabular_logger as tlogger\n'), ((94, 12, 94, 74), 'tabular_logger.exception', 'tlogger.exception', ({(94, 30, 94, 73): '"""WorkerHub._handle_output exception thrown"""'}, {}), "('WorkerHub._handle_output exception thrown')", True, 'import tabular_logger as tlogger\n'), ((134, 12, 134, 77), 'tabular_logger.exception', 'tlogger.exception', ({(134, 30, 134, 76): '"""AsyncTaskHub._handle_output exception thrown"""'}, {}), "('AsyncTaskHub._handle_output exception thrown')", True, 'import tabular_logger as tlogger\n'), ((67, 20, 67, 85), 'tabular_logger.info', 'tlogger.info', ({(67, 33, 67, 84): '"""WorkerHub._handle_input NO MORE WORKERS AWAILABLE"""'}, {}), "('WorkerHub._handle_input NO MORE WORKERS AWAILABLE')", True, 'import tabular_logger as tlogger\n'), ((73, 20, 73, 84), 'tabular_logger.info', 'tlogger.info', ({(73, 33, 73, 83): '"""WorkerHub._handle_input NO MORE INPUTS AWAILABLE"""'}, {}), "('WorkerHub._handle_input NO MORE INPUTS AWAILABLE')", True, 'import tabular_logger as tlogger\n'), ((90, 20, 90, 65), 'tabular_logger.info', 'tlogger.info', ({(90, 33, 90, 64): '"""WorkerHub._handle_output done"""'}, {}), "('WorkerHub._handle_output done')", True, 'import tabular_logger as tlogger\n'), ((130, 20, 130, 68), 'tabular_logger.info', 'tlogger.info', ({(130, 33, 130, 67): '"""AsyncTaskHub._handle_output done"""'}, {}), "('AsyncTaskHub._handle_output done')", True, 'import tabular_logger as tlogger\n')]