import logging
import time

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

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


logger = logging.getLogger(__name__)


def _now():
  return datetime.now(timezone.utc)


class DistributedTaskManagerStore(AbstractDistributedTaskManager):

  def __init__(self, store, task_data_type='distributed_task'):
    super().__init__()

    # We should maybe take care to clean up old tasks lest our table grow huge, but this isn't a good way of doing it
    # self.lock = RLock()
    # self.clean_up_thread = Thread(target=self._clean_up, daemon=True)
    # self.clean_up_thread.start()

    self._task_data_type = task_data_type
    self._store = store

    # We don't specify a validation format here, but the data type will be structured like this:
    # {
    #   'task_type_id': non_empty_str,
    #   'parent_task_run_id': nullable(uuid),
    #   'input': anything,
    #   'start': iso_datetime,
    #   'in_progress': iso_datetime,
    #   '[end]': iso_datetime,
    #   '[output]': anything,
    #   '[exception]': anything
    # }

    store.init_data_type(task_data_type)
    store.init_index(task_data_type, 'task_type_id', str)
    store.init_index(task_data_type, 'start', datetime)
    store.init_index(task_data_type, 'in_progress', datetime)
    store.init_index_getter(task_data_type, 'complete', lambda i: bool(i.get('end')), bool)

  def _put_task(self, task_type_id, params, parent_task_run_id):
    task = {
      'task_type_id': task_type_id,
      'parent_task_run_id': parent_task_run_id,
      'input': params,
      'start': _now(),
      'in_progress': datetime.utcfromtimestamp(0).replace(tzinfo=timezone.utc)
    }
    return self._store.create(self._task_data_type, task)

  def _update_task(self, task_run_id, to_update):
    task = self._store.get(self._task_data_type, task_run_id)
    task.update(to_update)
    self._store.update(self._task_data_type, task_run_id, task)

  def _put_task_in_progress(self, task_run_id):
    self._update_task(task_run_id, { 'in_progress': _now() })

  def _put_task_output(self, task_run_id, output):
    self._update_task(task_run_id, { 'end': _now(), 'output': output })

  def _put_task_exception(self, task_run_id, exception):
    self._update_task(task_run_id, { 'end': _now(), 'exception': exception })

  def _checkout_task(self, task_type_id):
    inactive_limit = _now() - timedelta(seconds=self.DISTRIBUTED_TASK_NO_ACTIVITY_TIMEOUT_S)
    compound_query = [ ('task_type_id', task_type_id), ('in_progress', '<', inactive_limit), ('complete', False) ]
    return self._store.index.find(self._task_data_type, compound_query, order_by='start')

  def _get_task_output(self, task_run_id):
    try:
      task = self._store.get(self._task_data_type, task_run_id)
    except NotFoundException as e:
      raise TaskNotFoundException() from e
    if not 'end' in task:
      raise TaskNotFinishedException()
    if 'output' in task:
      return task['output']
    raise TaskException(task['exception'])

  def _clean_up(self):
    # TODO: this is not currently used...
    # TODO: you should optimize this - no need to do a full table scan every time...
    while True:
      logger.debug("Cleaning up tasks...")
      with self.lock:
        to_delete = set()
        for task_run_id, task in self._store.list(self._task_data_type):
          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 task.get('end')) 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:
          self._store.delete(self._task_data_type, task_run_id)

      time.sleep(self.DISTRIBUTED_TASK_NO_ACTIVITY_TIMEOUT_S)
