from threading import RLock, Thread

from .distributed_task import DistributedTask


class AbstractDistributedTaskManager:

  DISTRIBUTED_TASK_NO_ACTIVITY_TIMEOUT_S = 3600
  DISTRIBUTED_TASK_MAX_AGE_S = 86400
  DISTRIBUTED_TASK_AWAIT_POLL_INTERVAL_S = 30
  DISTRIBUTED_TASK_WORKER_POLL_INTERVAL_S = 30
  DISTRIBUTED_TASK_WORKER_IN_PROGRESS_INTERVAL_S = 10
  DISTRIBUTED_TASK_WORKER_IN_PROGRESS_TIMEOUT_S = 15

  def __init__(self):
    self.registered_tasks = {}
    self._worker_threads = []
    self.registered_tasks_lock = RLock()

  # ----------------------------------------------------------------------------
  # Abstract methods

  def _put_task(self, task_type_id, input, parent_task_run_id):
    """
    Submits a new task request. Note that if the task_type_id, input, and parent_task_run_id all match those of an
    existing task, the task run ID of that existing task will be returned.

    Returns a task run ID
    """

    raise NotImplementedError()

  def _put_task_in_progress(self, task_run_id):
    """
    Report that the specified task is in progress.
    """

    raise NotImplementedError()

  def _put_task_output(self, task_run_id, output):
    raise NotImplementedError()

  def _put_task_exception(self, task_run_id, exception):
    raise NotImplementedError()

  def _checkout_task(self, task_type_id):
    """
    Returns the next available task for this task_type_id, or None if there aren't any
    """

    raise NotImplementedError()

  def _get_task_output(self, task_run_id):
    """
    Returns the task output for the specified 'task_run_id' in the normal case. If the task completed with an exception,
    the message of that exception will be raised here (as a 'TaskException'). Note that this exception will not include
    the offending stack trace, as it may not even exist on this host.

    Raises 'TaskNotFoundException' if the 'task_run_id' could not be found.
    Raises 'TaskNotFinishedException' if the task has not finished.
    """

    raise NotImplementedError()

  def _clean_up(self):
    raise NotImplementedError()

  # ----------------------------------------------------------------------------

  def distributed_task(self, task_type_id, n_worker_threads=1):
    """
    A decorator factory that returns a DistributedTask instance.

    manager = LocalDistributedTaskManager()

    @manager.distributed_task("sum:1")
    def sum(a, b):
      return a + b

    # The decorator transforms 'sum' into a DistributedTask instance. These instances are callable, but bear in mind
    # that they could be slow depending on your manager's implementation.

    # This starts a new distributed task, gets picked up by a worker, and eventually makes the output available:
    result = sum(1, 2)

    # To start the distributed task asynchronously:
    task_run_id = sum.start(1, 2)

    # Then later you can wait on the output:
    result = sum.wait(task_run_id)
    """

    # TODO: derive the 'task_type_id' automatically based on the function body

    def decorator(f):
      with self.registered_tasks_lock:
        if task_type_id in self.registered_tasks:
          raise Exception("Task type already registered: %s" % task_type_id)
        task = DistributedTask(self, task_type_id, f)
        self.registered_tasks[task_type_id] = task
        for _ in range(n_worker_threads):
          thread = Thread(target=task.worker, daemon=True)
          thread.start()
          self._worker_threads.append(thread)
        return task
    return decorator

  def join(self):
    for thread in self._worker_threads:
      thread.join()




class AbstractTaskStore:

  def create_task(self, task_type_id, parent_task_run_id, pickled_input):
    '''
    Returns the task ID
    '''
    raise NotImplementedError()

  def update_task_in_progress(self, task_run_id):
    raise NotImplementedError()

  def update_task_output(self, task_run_id, pickled_output):
    raise NotImplementedError()

  def update_task_exception(self, task_run_id, pickled_exception):
    raise NotImplementedError()

  def checkout_task(self, task_type_id):
    """
    Returns the next available task for this task_type_id, or None if there aren't any
    """

    raise NotImplementedError()

  def get_task_output(self, task_run_id):
    """
    Returns the task output for the specified 'task_run_id' in the normal case. If the task completed with an exception,
    the message of that exception will be raised here (as a 'TaskException'). Note that this exception will not include
    the offending stack trace, as it may not even exist on this host.

    Raises 'TaskNotFoundException' if the 'task_run_id' could not be found.
    Raises 'TaskNotFinishedException' if the task has not finished.
    """

    raise NotImplementedError()


class Task:

  def __init__(self, manager, key, n_worker_threads):
    self.manager = manager
    self.key = key
    self.n_worker_threads = n_worker_threads

  def async(self, *args, **kwargs):
    task_run_id = self.manager.task_store.create_task(self.key, parent_key, pickle.dumps((args, kwargs)))

    # Notify our worker to check for new tasks right away - this makes it likely that the same host that started the
    # task will also process it (and avoids having to wait for the DISTRIBUTED_TASK_WORKER_POLL_INTERVAL_S)
    self.worker_notify_event.set()

    return task_run_id

  def __call__(self, *args, **kwargs):
    task_run_id = self.async(*args, **kwargs)
    return self.manager.wait(task_run_id)

  def terminate_workers(self):
    self.__can_work = False
    self.worker_notify_event.set()

  def worker(self):
    self.__can_work = True
    while self.__can_work:
      logger.debug("Looking for '%s' tasks..." % self.task_type_id)

      # Get the next available task for this task_type_id
      response = self.manager._checkout_task(self.task_type_id)
      if response:
        task_run_id, task = response
        # Start a thread to report that we are working on this task
        is_complete_event = Event()
        def update_last_in_progress():
          with suppress(TaskFinishedException):
            while not is_complete_event.wait(self.manager.DISTRIBUTED_TASK_WORKER_IN_PROGRESS_INTERVAL_S):
              logger.debug("Reporting '%s' task in progress with ID: %s" % (self.task_type_id, task_run_id))
              self.manager._put_task_in_progress(task_run_id)
        in_progress_thread = Thread(target=update_last_in_progress, daemon=True)
        in_progress_thread.start()

        if not hasattr(DistributedTask.thread_local_storage, 'task_run_id_stack'):
          DistributedTask.thread_local_storage.task_run_id_stack = []
        DistributedTask.thread_local_storage.task_run_id_stack.append(task_run_id)

        logger.debug("Handling '%s' task with ID: %s (parent task: %s)" % (self.task_type_id, task_run_id, task['parent_task_run_id']))

        try:
          args, kwargs = task['input']
          output = self.handler(*args, **kwargs)
          self.manager._put_task_output(task_run_id, output)
        except Exception as e:
          logger.exception("Distributed task worker encountered exception")
          self.manager._put_task_exception(task_run_id, e)

        logger.debug("Done handling '%s' task with ID: %s (parent task: %s)" % (self.task_type_id, task_run_id, task['parent_task_run_id']))

        self.thread_local_storage.task_run_id_stack.pop()

        is_complete_event.set()
        in_progress_thread.join()
        self.await_notify_event.set()

        # Check right away for more tasks
        continue

      if self.worker_notify_event.wait(self.manager.DISTRIBUTED_TASK_WORKER_POLL_INTERVAL_S):
        self.worker_notify_event.clear()

    logger.info("Worker for '%s' task is exiting" % self.task_type_id)


DISTRIBUTED_TASK_NO_ACTIVITY_TIMEOUT_S = 3600
DISTRIBUTED_TASK_MAX_AGE_S = 86400
DISTRIBUTED_TASK_AWAIT_POLL_INTERVAL_S = 30
DISTRIBUTED_TASK_WORKER_POLL_INTERVAL_S = 30
DISTRIBUTED_TASK_WORKER_IN_PROGRESS_INTERVAL_S = 10
DISTRIBUTED_TASK_WORKER_IN_PROGRESS_TIMEOUT_S = 15


class Manager:

  def __init__(self, task_store):
    self._task_store = task_store
    self._registered_tasks = {}
    self._await_notify_event = Event()
    self._worker_notify_event = Event()


  def wait(self, task_run_id):
    while True:
      with suppress(TaskNotFinishedException):
        return self._task_store.get_task_output(task_run_id)
      if self._await_notify_event.wait(DISTRIBUTED_TASK_AWAIT_POLL_INTERVAL_S):
        self._await_notify_event.clear()


  def task(key=None, n_worker_threads=1):

    if not key:
      # Derive the key automatically by hashing the function code

    def decorator(f):
      with self.registered_tasks_lock:
        if task_type_id in self._registered_tasks:
          raise Exception("Task type already registered: %s" % task_type_id)
        task = DistributedTask(self, task_type_id, f)
        self._registered_tasks[task_type_id] = task
        for _ in range(n_worker_threads):
          thread = Thread(target=task.worker, daemon=True)
          thread.start()
          self._worker_threads.append(thread)
        return task

    return decorator





manager = Manager(task_store)

@manager.task(key=None, n_worker_threads=1)
def sum(a, b):
  return a + b


id = sum.async(a, b)
manager.wait(id)









