import logging
import time

from datetime import datetime, timezone
from threading import RLock, Thread
from uuid import uuid4

from ..abstract_distributed_task_manager import AbstractDistributedTaskManager
from ..exceptions import *


logger = logging.getLogger(__name__)


class DistributedTaskManagerSingleProcess(AbstractDistributedTaskManager):

  def __init__(self):
    super().__init__()
    self.tasks = {}
    self.lock = RLock()

    self.clean_up_thread = Thread(target=self._clean_up, daemon=True)
    self.clean_up_thread.start()

  def _is_started(self, task):
    return ('in_progress' in task)

  def _is_finished(self, task):
    return ('end' in task)

  def _is_in_progress(self, task):
    if not self._is_started(task):
      return False
    if self._is_finished(task):
      return False
    now = datetime.now(timezone.utc)
    return ((now - task['in_progress']).total_seconds() < self.DISTRIBUTED_TASK_WORKER_IN_PROGRESS_TIMEOUT_S)

  def _is_expired(self, task):
    if not self._is_started(task):
      return False
    if self._is_finished(task):
      return False
    now = datetime.now(timezone.utc)
    return ((now - task['in_progress']).total_seconds() < self.DISTRIBUTED_TASK_WORKER_IN_PROGRESS_TIMEOUT_S)

  def _clean_up(self):
    while True:
      logger.debug("Cleaning up tasks...")
      with self.lock:
        now = datetime.now(timezone.utc)
        to_delete = set()
        for task_run_id, task in self.tasks.items():
          last_activity = task.get('end') or task.get('in_progress') or task['start']
          activitiy_age_s = (now - last_activity).total_seconds()
          if (activitiy_age_s > self.DISTRIBUTED_TASK_MAX_AGE_S):
            logger.debug("Removing task record due to age: %s" % task)
            to_delete.add(task_run_id)
            continue
          if (not self._is_finished(task)) and (activitiy_age_s > self.DISTRIBUTED_TASK_NO_ACTIVITY_TIMEOUT_S):
            msg = "No activity for task in %is (marking it as failed): %s" % (activitiy_age_s, task)
            e = NoActivityException(msg)
            logger.error(msg)
            self._put_task_exception(task_run_id, e)
            continue
        for task_run_id in to_delete:
          del self.tasks[task_run_id]

      time.sleep(self.DISTRIBUTED_TASK_NO_ACTIVITY_TIMEOUT_S)

  def _put_task(self, task_type_id, input, parent_task_run_id):
    with self.lock:
      if parent_task_run_id:
        for task in self.tasks.values():
          if (task['task_type_id'] == task_type_id):
            if (task['input'] == input) and (task['parent_task_run_id'] == parent_task_run_id):
              return task['task_run_id']

      task_run_id = str(uuid4())
      self.tasks[task_run_id] = {
        'task_type_id': task_type_id,
        'task_run_id': task_run_id,
        'start': datetime.now(timezone.utc),
        'input': input,
        'parent_task_run_id': parent_task_run_id
      }
      return task_run_id

  def _put_task_in_progress(self, task_run_id):
    with self.lock:
      task = self._get_task(task_run_id)
      if self._is_finished(task):
        raise TaskFinishedException()
      task['in_progress'] = datetime.now(timezone.utc)
      self.tasks[task_run_id] = task

  def _put_task_output(self, task_run_id, output):
    with self.lock:
      task = self._get_task(task_run_id)
      task['finished_at'] = datetime.now(timezone.utc)
      task['output'] = output
      self.tasks[task_run_id] = task

  def _put_task_exception(self, task_run_id, exception):
    with self.lock:
      task = self._get_task(task_run_id)
      task['finished_at'] = datetime.now(timezone.utc)
      task['exception'] = str(exception)
      self.tasks[task_run_id] = task

  def _checkout_task(self, task_type_id):
    with self.lock:
      oldest_task = None
      for task in self.tasks.values():
        if (task['task_type_id'] == task_type_id):
          if (not self._is_finished(task)) and (not self._is_in_progress(task)):
            if (not oldest_task) or (task['start'] < oldest_task['start']):
              oldest_task = task
      if oldest_task:
        self._put_task_in_progress(oldest_task['task_run_id'])
      return oldest_task

  def _get_task(self, task_run_id):
    with self.lock:
      task = self.tasks.get(task_run_id, None)
      if not task:
        raise TaskNotFoundException()
      return task

  def _get_child_tasks(self, task_run_id):
    with self.lock:
      child_tasks = []
      for task in self.tasks.values():
        if (task['parent_task_run_id'] == task_run_id):
          child_tasks.append(task)
      return child_tasks

  def _get_task_output(self, task_run_id):
    with self.lock:
      task = self._get_task(task_run_id)
      if not self._is_finished(task):
        raise TaskNotFinishedException()
      if 'output' in task:
        return task['output']
      raise TaskException("Distributed task raised exception (trace can be found on the worker): " + task['exception'])
