import os
from .worker_manager import ThreadWorkerQueue
import logging

log = logging.getLogger().getChild("async_tasks_manager")

THREAD_COUNT = int(os.getenv("ASYNC_WORKERS", 4))
IO_THREAD_COUNT = int(os.getenv("IO_WORKERS", 50))
ThreadWorkerQueue.instantiate(thread_count=THREAD_COUNT, io_worker_limit=IO_THREAD_COUNT)

acceptable_async_sources = {}
# just remap the static function to module
io_task = ThreadWorkerQueue.io_task
io_task_limited = ThreadWorkerQueue.io_task_limited


def async_task(_func=None, *, async_source=None):
    """Decorates async task. use event_source parameter to indicate,
    that task is event handler for particularly named source events."""
    def decorator(func):
        def default_wrap(*args, task_handle=None, **kwargs):
            try:
                for arg in args:
                    kwargs.update(arg)
                log.info(f"ASYNC_ENDPOINT: {async_source}: {kwargs}")
                return func(*args, **kwargs, task_handle=task_handle)
            except Exception as e:
                log.exception(e)
                if not task_handle.immediate_response_sent():
                    task_handle.send_immediate_error(e)
        ref = ThreadWorkerQueue.task(default_wrap)
        acceptable_async_sources[async_source] = ref
        return ref

    if _func is None:
        return decorator
    else:
        return ThreadWorkerQueue.task(_func)


def knows_how_to_handle(async_data):
    """Take async calls, that have arrived from outside sources

    :param async_data: dict, that has to have{"source": "name"}, requires a handler definition of

            @async_task_manager.async_task(event_source="name")
            def event_handler(event_data, task_handle):
                task_handle.send_immediate_response('status':'accepted')
                return {'result':'stuff'}

    to succeed. task_handle.send_immediate_response needs to be called ASAP, so the async call result can be returned.
    """
    source = async_data.get('async_source', None)
    if source in acceptable_async_sources:
        # fire off the async handler
        return acceptable_async_sources[source](async_data)
    return None