'''
TODO:
- Figure out how to derive the task_type_id automatically
- Saving work from parent / child tasks?

'''


import boto3
import logging
import sched
import time

from concurrent.futures import ThreadPoolExecutor
from contextlib import suppress
from datetime import datetime, timedelta, timezone
from op_storage import Condition
from threading import Event, Lock, Thread
from uuid import UUID


logger = logging.getLogger(__name__)


 # There is an exponential backoff so that we're not furiously polling the DB for long-running tasks
MIN_WAIT_POLL_INTERVAL_S = 1
MAX_WAIT_POLL_INTERVAL_S = 300

# MIN_HEARTBEAT_TIMEOUT_S = 270
# MAX_HEARTBEAT_TIMEOUT_S = 1800

# We want to set the deadline for a task to exactly when the next heartbeat should occur (which in turn updates that
# deadline), but the timing won't be exact. Adding this buffer helps ensure that we will never encounter
# properly-running tasks that have missed their deadline.
HEARTBEAT_BUFFER_S = 10

# # The manager polls the DB looking for abandoned tasks this often - note that this must be longer than the
# # 'MIN_HEARTBEAT_TIMEOUT_S' interval, but there's no sense in making it too much longer.
# ABANDONED_POLL_INTERVAL_S = 300




MIN_HEARTBEAT_TIMEOUT_S = 1
MAX_HEARTBEAT_TIMEOUT_S = 1800



# The manager polls the DB looking for abandoned tasks this often - note that this must be longer than the
# 'MIN_HEARTBEAT_TIMEOUT_S' interval, but there's no sense in making it too much longer.
ABANDONED_POLL_INTERVAL_S = 2



class OPTaskException(Exception):
  pass


class TaskException(OPTaskException):
  pass


class TaskFinishedException(OPTaskException):
  pass


class TaskNotFinishedException(OPTaskException):
  pass


class TaskNotFoundException(OPTaskException):
  pass


class NoActivityException(OPTaskException):
  pass


class TaskWorker:

  def __init__(self, manager, task_type_id, n_workers, fn):
    self._m = manager
    self._task_type_id = task_type_id
    self._n_workers = n_workers
    self._fn = fn
    self._thread_pool_executor = ThreadPoolExecutor(max_workers=n_workers)
    self._heartbeat_thread_pool_executor = ThreadPoolExecutor(max_workers=n_workers)


  def _submit(self, task_run_id, task):
    self._thread_pool_executor.submit(self._run, task_run_id, task, Lock(), Event())


  def _start_heartbeat(self, delay_s, task_run_id, task, lock, event):
    self._heartbeat_thread_pool_executor.submit(self._heartbeat, delay_s, task_run_id, task, lock, event)


  def _heartbeat(self, delay_s, task_run_id, task, lock, event):
    if event.wait(delay_s):
      event.clear()
    with lock:
      logger.debug('Reporting heartbeat for task run ID: %s' % task_run_id)
      print('Reporting heartbeat for task run ID: %s' % task_run_id)
      try:
        if task.get('end'):
          return
        i = task['n_beats']
        next_delay_s = _get_exp(i, MIN_HEARTBEAT_TIMEOUT_S, MAX_WAIT_POLL_INTERVAL_S)
        task['deadline'] = datetime.now(timezone.utc) + timedelta(seconds=next_delay_s + HEARTBEAT_BUFFER_S)
        task['n_beats'] = i + 1
        self._m._db.update(self._m._data_type, task_run_id, task)
      except:
        msg = 'Exception reporting heartbeat'
        logger.exception(msg)
        self._m._op_alarm and self._m._op_alarm.report_issue(msg, metadata=self._m._service_metadata)
      finally:
        self._start_heartbeat(next_delay_s, task_run_id, task, lock, event)


  def _run(self, task_run_id, task, lock, event):
    self._start_heartbeat(MIN_HEARTBEAT_TIMEOUT_S, task_run_id, task, lock, event)
    try:
      args, kwargs = task['input']
      task['output'] = self._fn(*args, **kwargs)
    except Exception as e:
      logger.exception('Distributed task worker encountered exception')
      task['exc'] = e
    with lock:
      task['deadline'] = None
      task['end'] = datetime.now(timezone.utc)
      self._m._db.update(self._m._data_type, task_run_id, task)
    event.set()


  def dispatch(self, *args, **kwargs):
    now = datetime.now(timezone.utc)
    task = {
      'svc': self._m._service_metadata,
      'task_type_id': self._task_type_id,
      'parent_task_run_id': None,
      'input': (args, kwargs),
      'start': now,
      'n_beats': 0,
      'deadline': now + timedelta(seconds=MIN_HEARTBEAT_TIMEOUT_S + HEARTBEAT_BUFFER_S)
    }
    task_run_id = self._m._db.create(self._m._data_type, task)

    # Only push the message to SQS if our own '_thread_pool_executor' is full. Otherwise, we can save some overhead and
    # just submit it directly.
    # Warning: we can only determine the number of tasks in progress via an ostensibly private call
    if self._m._sqs and (self._thread_pool_executor._work_queue.qsize() >= self._n_workers):
      self._m._sqs.send_message(QueueUrl=self._m._queue_url, MessageBody=str(task_run_id))
    else:
      self._submit(task_run_id, task)

    return task_run_id


  def __call__(self, *args, **kwargs):
    task_run_id = self.dispatch(*args, **kwargs)
    return self._m.wait(task_run_id)


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

  #     # Get the next available task for this task_type_id
  #     response = self.m._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)






def _get_exp(i, min_value, max_value):
  return min((2 ^ i) * min_value, max_value)


class Manager:


  def __init__(self, service_metadata, op_storage, op_alarm=None, aws_config=None, queue_name='op_distributed_task', data_type='op_distributed_task'):
    self._service_metadata = service_metadata

    self._data_type = data_type
    self._db = op_storage
    op_storage.init(data_type, indexes=[ 'deadline' ])

    self._op_alarm = op_alarm

    self._sqs = None
    if aws_config:
      self._sqs = boto3.client('sqs', **aws_config)
      response = self._sqs.create_queue(QueueName=queue_name)
      self._queue_url = response['QueueUrl']

    self._registered_tasks = {}
    self._registered_tasks_lock = Lock()

    self._abandoned_event = Event()
    self._wait_event = Event()


  def start(self):
    if self._sqs:
      Thread(target=self._poll_sqs_tasks, daemon=True).start()
    Thread(target=self._poll_abandoned_tasks, daemon=True).start()


  def _poll_sqs_tasks(self):
    while True:
      try:
        response = self._sqs.receive_message(QueueUrl=self._queue_url, WaitTimeSeconds=20)
        for msg in response.get('Messages', []):
          task_run_id = UUID(msg['Body'])
          task = self._db.get(self._data_type, task_run_id)
          if task['task_type_id'] not in self._registered_tasks:
            # We don't know how to handle this message - just leave it in the queue
            continue
          # Remove this message from the queue and push this message to the appropriate worker
          self._sqs.delete_message(QueueUrl=self._queue_url, ReceiptHandle=msg['ReceiptHandle'])
          task_worker = self._registered_tasks[task['task_type_id']]
          task_worker._submit(task_run_id, task)
      except:
        msg = 'Exception handling SQS tasks'
        logger.exception(msg)
        self._op_alarm and self._op_alarm.report_issue(msg, metadata=self._service_metadata)


  def _poll_abandoned_tasks(self):
    while True:
      if self._abandoned_event.wait(ABANDONED_POLL_INTERVAL_S):
        self._abandoned_event.clear()
      try:
        now = datetime.now(timezone.utc)
        for task_run_id, task in self._db.iter(self._data_type, Condition('deadline') < now):
          if task['task_type_id'] in self._registered_tasks:
            task['deadline'] = datetime.now(timezone.utc) + timedelta(seconds=MIN_HEARTBEAT_TIMEOUT_S + HEARTBEAT_BUFFER_S)
            self._db.update(self._data_type, task_run_id, task)
            task_worker = self._registered_tasks[task['task_type_id']]
            task_worker._submit(task_run_id, task)
            continue
          # We found an abandoned task that we don't know how to complete
          if (task['deadline'] + timedelta(seconds=ABANDONED_POLL_INTERVAL_S)) > now:
            # There could be other workers who DO know how to handle that task type at this same point in the loop, so
            # let's give them a chance to update / finish the task before we alarm on it.
            continue
          msg = (
            'Found abandoned task with unrecognized type - you need to spin up a worker that knows how to handle '
            'this task type. You can do so by running the following command: op-dev run --pkg-id %(pkg_id)s '
            '--commit-id %(commit_id)s --config-id %(config_id)s'
          ) % task['svc']
          logger.critical(msg)
          self._op_alarm and self._op_alarm.report_issue(msg, task=task, metadata=self._service_metadata)
      except:
        msg = 'Exception handling abandoned tasks'
        logger.exception(msg)
        self._op_alarm and self._op_alarm.report_issue(msg, metadata=self._service_metadata)


  def __call__(self, task_type_id, n_workers=1):
    # TODO: make 'task_type_id' optional
    if not task_type_id:
      # Derive the task_type_id automatically using the commit / config?
      # Or maybe hash the funciton code is better? Eh, not necessarily unless it's a true global hash...
      raise Exception()

    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_worker = TaskWorker(self, task_type_id, n_workers, f)
        self._registered_tasks[task_type_id] = task_worker
      return task_worker
    return decorator


  def wait(self, task_run_id):
    i = 0
    while True:
      wait_s = _get_exp(i, MIN_WAIT_POLL_INTERVAL_S, MAX_WAIT_POLL_INTERVAL_S)
      if self._wait_event.wait(wait_s):
        self._wait_event.clear()
      with suppress(TaskNotFinishedException):
        task = self._db.get(self._data_type, task_run_id)
        if 'output' in task:
          return task['output']
        if 'exc' in task:
          raise TaskException(task['exc'])
      i += 1







# TEST
from op_storage import factory
from uuid import UUID


# Must have administrative access for SQS
aws_config = {
  "aws_access_key_id": "AKIA42TXSCQG2S26B6GJ",
  "aws_secret_access_key": "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",
  "region_name": "us-west-2"
}
db_config = {
  "connection_string": "mongodb+srv://pro-ipsi:OfWxgYawD1XUwvj6@cluster-0-hpws1.mongodb.net/?retryWrites=true&w=majority&ssl=true&tlsAllowInvalidCertificates=true",
  "database": "pro-ipsi"
}
db = factory.build('mongo_db', db_config)

# Note that you still have to 'start' the manager
task = Manager(aws_config, db)

# This starts the worker thread (or multiple if specified) looking for new tasks
@task('1')
def add(a, b):
  print('Running task...')
  import time
  time.sleep(5)
  return a + b

task.start()

# print('Dispatching new task')
# task_run_id = add.dispatch(4, 5)

task_run_id = UUID('b5893643-a081-41c5-97a6-5d562d342c25')

print('Waiting on task:')
result = task.wait(task_run_id)

print('Got result:')
print(result)




