import logging
import threading

from contextlib import suppress
from threading import Event, Thread

from .exceptions import *


logger = logging.getLogger(__name__)


def _pretty_str_task(task):
  output = "type=%(task_type_id)s id=%(task_run_id)s input=%(input)s ( %(start)s - " % task
  output += task['in_progress'].isoformat() if 'in_progress' in task else 'None'
  output += " - "
  output += task['end'].isoformat() if 'end' in task else 'None'
  return output + " )"


def _pretty_print_task_tree(task_tree, _indent=0):
  print(("  " * _indent) + _pretty_str_task(task_tree['task']))
  for sub_task in task_tree['sub_tasks']:
    _pretty_print_task_tree(sub_task, _indent=(_indent + 2))


class DistributedTask:

  thread_local_storage = threading.local()

  def __init__(self, manager, task_type_id, handler):
    self.manager = manager
    self.task_type_id = task_type_id
    self.handler = handler

    self.await_notify_event = Event()
    self.worker_notify_event = Event()

  def __call__(self, *args, **kwargs):
    task_run_id = self.start(*args, **kwargs)
    output = self.wait(task_run_id)
    return output

  def start(self, *args, **kwargs):
    parent_task_run_id = None
    with suppress(Exception):
      parent_task_run_id = DistributedTask.thread_local_storage.task_run_id_stack[-1]
    task_run_id = self.manager._put_task(self.task_type_id, (args, kwargs), parent_task_run_id)

    # Note: just because we submitted a new task does not necessarily mean it will run - if there is a
    # parent_task_run_id and there is already a pending task matching the parameters of this one, then we will just
    # return the ID of that pending task.
    msg = "Submitted new '%s' task with ID: %s" % (self.task_type_id, task_run_id)
    if parent_task_run_id:
      msg += " (parent task ID: %s)" % parent_task_run_id
    logger.info(msg)

    # 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 wait(self, task_run_id):
    while True:
      logger.debug("Looking for results for task ID: %s" % task_run_id)

      with suppress(TaskNotFinishedException):
        return self.manager._get_task_output(task_run_id)
      if self.await_notify_event.wait(self.manager.DISTRIBUTED_TASK_AWAIT_POLL_INTERVAL_S):
        self.await_notify_event.clear()

  # def get_task(self, task_run_id):
  #   return self.manager._get_task(task_run_id)

  # def get_child_tasks(self, task_run_id):
  #   return self.manager._get_child_tasks(task_run_id)

  # def get_task_tree(self, task_run_id):
  #   root = { 'task': self.get_task(task_run_id), 'sub_tasks': [] }
  #   for child_task_run_id, child_task in self.get_child_tasks(task_run_id):
  #     root['sub_tasks'].append(self.get_task_tree(child_task_run_id))
  #   return root

  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)
