'''
Defines a TaskManager build on top of an op_storage DB instance

TODO:
- Documentation
- I'm worried about going overboard with threads - how about we just report the heartbeats from a single thread?



Simplify submission to the thread pool executors
https://www.aeracode.org/2018/02/19/python-async-simplified/
def get_chat_id(name):
    return Chat.objects.get(name=name).id

async def main():
    executor = concurrent.futures.ThreadPoolExecutor(max_workers=3)
    loop = asyncio.get_event_loop()
    result = await loop.run_in_executor(executor, get_chat_id, "django")
'''

import asyncio
import atexit
import inspect
import logging
import pickle
import sys
import threading
import time
import traceback

from concurrent.futures import ThreadPoolExecutor
from concurrent.futures.thread import _python_exit as _concurrent_futures_thread_python_exit
from contextlib import contextmanager, suppress
from datetime import datetime, timedelta, timezone
from op_storage import Field
from op_storage.exceptions import NotFoundException
from threading import Event, Lock, RLock, Thread

from .env import OP_DEV_METADATA

# Unregister the exit handler associated with the ThreadPoolExecutor - the entire purpose of our library is to handle
# interrupted tasks, so we want those threads to behave like true daemons.
# See this for more info:
# https://stackoverflow.com/questions/49992329/the-workers-in-threadpoolexecutor-is-not-really-daemon
# WARNING: this changes the default behavior of 'concurrent.futures.thread' (for all consumers)
atexit.unregister(_concurrent_futures_thread_python_exit)


logger = logging.getLogger(__name__)


MIN_HEARTBEAT_S = 25
MAX_HEARTBEAT_S = 600
HEARTBEAT_EXP_GROWTH_FACTOR = 10

# 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.
ABANDONED_POLL_INTERVAL_S = 30


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


def _exp(i, growth_factor, min_value, max_value):
  # For computing exponential growth from the 'min_value' up to the 'max_value'
  return min(((2 ** i) -1) * growth_factor + min_value, max_value)


def _build_abandoned_task_msg(task):
  msg = (
    'Found abandoned task with unrecognized type - you need to spin up a worker that knows how to handle this task '
    'type'
  )
  task_metadata = task.get('source')
  if task_metadata and all(task_metadata.get(k) for k in ('pkg_id', 'commit_id', 'config_id')):
    msg += (
      '. 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_metadata
  else:
    msg += ': ' + task['task_type_id']
  return msg


def _generate_task_type_id(fn):
  # Derives a task type ID from the given function automatically
  return f'{fn.__name__}{inspect.signature(fn)}'


class OPTaskException(Exception):
  pass


class TaskResult:

  def __init__(self, task):
    self._task = task

  def result(self):
    if 'output' in self._task:
      return self._task['output']
    raise OPTaskException(self._task['exc'])


class LocalTaskResult:

  def __init__(self, task, task_completed_event):
    self._task = task
    self._task_completed_event = task_completed_event

  def result(self):
    self._task_completed_event.wait()
    if 'output' in self._task:
      return self._task['output']
    raise OPTaskException(self._task['exc'])


class DistributedTaskResult:

  def __init__(self, task_run_id, manager, heartbeat_s):
    self._task_run_id = task_run_id
    self._m = manager
    self._heartbeat_s = heartbeat_s # Potentially None

  def result(self):
    i = 0
    while True:
      # We base our polling frequency on the heartbeat
      wait_s = self._heartbeat_s or _exp(i, HEARTBEAT_EXP_GROWTH_FACTOR, MIN_HEARTBEAT_S, MAX_HEARTBEAT_S)
      time.sleep(wait_s)
      try:
        task = self._m._service.db.get(self._m._data_type, self._task_run_id)
      except Exception as e:
        # Just adding amore descriptive error message
        raise Exception('Failed polling for distributed task (the underlying task may have succeeded)') from e
      if 'output' in task:
        return task['output']
      if 'exc' in task:
        raise OPTaskException(task['exc'])
      i = task['n_beats']


class TaskWorker:

  def __init__(self, manager, task_type_id, max_workers, heartbeat_s, fn):
    self._m = manager
    self._task_type_id = task_type_id
    self._heartbeat_s = heartbeat_s # Potentially None
    self._fn = fn
    self._thread_pool_executor = ThreadPoolExecutor(max_workers=max_workers)
    self._heartbeat_thread_pool_executor = ThreadPoolExecutor(max_workers=max_workers)
    self._n_in_progress = 0
    self._n_in_progress_lock = RLock()


  def _schedule_heartbeat(self, delay_s, task_run_id, task, lock, event):
    logger.debug('Scheduling heartbeat with %ss delay for task run ID: %s' % (delay_s, task_run_id))
    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, heartbeat_lock, task_completed_event):
    if task_completed_event.wait(delay_s):
      return
    with heartbeat_lock:
      with self._m._exc_handler('Exception reporting heartbeat'):
        if task.get('end'):
          return
        logger.debug('Reporting heartbeat for task run ID: %s' % task_run_id)
        i = task['n_beats']
        next_delay_s = self._heartbeat_s or _exp(i, HEARTBEAT_EXP_GROWTH_FACTOR, MIN_HEARTBEAT_S, MAX_HEARTBEAT_S)
        task['deadline'] = _now() + timedelta(seconds=next_delay_s + HEARTBEAT_BUFFER_S)
        task['n_beats'] = i + 1
        self._m._service.db.update(self._m._data_type, task_run_id, task)
      self._schedule_heartbeat(next_delay_s, task_run_id, task, heartbeat_lock, task_completed_event)


  def _run(self, task_run_id, task, heartbeat_lock, task_completed_event):
    logger.debug("Beginning work on task run ID '%s' for task type ID '%s'" % (task_run_id, self._task_type_id))
    heartbeat_delay_s = self._heartbeat_s or MIN_HEARTBEAT_S
    task['n_beats'] = 0
    task['deadline'] = _now() + timedelta(seconds=heartbeat_delay_s + HEARTBEAT_BUFFER_S)
    self._schedule_heartbeat(heartbeat_delay_s, task_run_id, task, heartbeat_lock, task_completed_event)
    try:
      with self._m._parent(task_run_id):
        args, kwargs = pickle.loads(task['p_input'])
        if asyncio.iscoroutinefunction(self._fn):
          task['output'] = asyncio.run(self._fn(*args, **kwargs))
        else:
          task['output'] = self._fn(*args, **kwargs)
    except Exception as e:
      task['exc'] = e
      msg = "Exception for task type ID '%s'" % self._task_type_id
      logger.exception(msg)
      self._m._report_task_exception(msg, task_run_id, task)
    with heartbeat_lock:
      task['deadline'] = None
      task['end'] = _now()
      with self._m._exc_handler('Exception reporting task completion'):
        # If this call fails, the task will appear to be abandoned and get re-run at some point
        self._m._service.db.update(self._m._data_type, task_run_id, task)
    task_completed_event.set()
    with self._n_in_progress_lock:
      self._n_in_progress -= 1
    logger.debug("Finished work on task run ID '%s' for task type ID '%s'" % (task_run_id, self._task_type_id))


  def _execute_task(self, task_run_id, task):
    with self._n_in_progress_lock:
      self._n_in_progress += 1
      task_completed_event = Event()
      self._thread_pool_executor.submit(self._run, task_run_id, task, Lock(), task_completed_event)
    return LocalTaskResult(task, task_completed_event)


  @contextmanager
  def _get_capacity(self):
    with self._n_in_progress_lock:
      yield max(self._thread_pool_executor._max_workers - self._n_in_progress, 0)


  # TODO: expand this interface to allow for scheduled tasks where 't' is a timezone-aware datetime
  # def schedule(self, t, *args, **kwargs):
  #   fn.schedule(<datetime>, *args, **kwargs)


  # TODO: expose scheduled / compelted / errored tasks from the db via an API


  def submit(self, *args, **kwargs):
    parent_task_run_id = self._m._get_parent()
    logger.debug("Submitting task of type '%s' with parent '%s'" % (self._task_type_id, parent_task_run_id))
    pickled_input = pickle.dumps((args, kwargs))
    if parent_task_run_id:
      results = self._m._db.list(
        self._m._data_type,
        Field('parent') == parent_task_run_id,
        Field('task_type_id') == self._task_type_id,
        Field('p_input') == pickled_input
      )
      if results:
        task_run_id, task = results[0]
        if 'end' in task:
          logger.debug('Found completed child task for parent: %s' % parent_task_run_id)
          return TaskResult(task)
        else:
          logger.debug('Found task in progress for parent: %s' % parent_task_run_id)
          return DistributedTaskResult(task_run_id, self._m, self._heartbeat_s)

    new_task = {
      'source': OP_DEV_METADATA,
      'task_type_id': self._task_type_id,
      'parent': parent_task_run_id,
      'p_input': pickled_input,
      'start': _now(),
      'n_beats': 0,
      'deadline': _now() + timedelta(seconds=(self._heartbeat_s or MIN_HEARTBEAT_S) + HEARTBEAT_BUFFER_S)
    }
    with self._get_capacity() as capacity:
      if capacity > 0:
        task_run_id = self._m._service.db.create(self._m._data_type, new_task)
        logger.debug('Executing task for task type ID: %s' % self._task_type_id)
        return self._execute_task(task_run_id, new_task)
      else:
        # We didn't have room for this task, so we deliberately abandon it so it can get picked up by someone else ASAP
        new_task['deadline'] = _now()
        task_run_id = self._m._service.db.create(self._m._data_type, new_task)
        logger.debug('Deferring execution of new task for task type ID: %s' % self._task_type_id)
        return DistributedTaskResult(task_run_id, self._m, self._heartbeat_s)


  def __call__(self, *args, **kwargs):
    return self.submit(*args, **kwargs).result()


def _try_str(value):
  with suppress(Exception):
    return str(args)
  return '?'


# Builds a 'task' for input to the 'report_issue' alarm function
def _build_issue_task(task_run_id, task):
  args, kwargs = pickle.loads(task['p_input'])
  return {
    'task_run_id': task_run_id,
    'source': task['source'],
    'task_type_id': task['task_type_id'],
    'start': task['start'],
    'n_beats': task['n_beats'],
    'deadline': task['deadline'],
    'input': {
      'args': [ _try_str(x) for x in args ],
      'kwargs': { x: _try_str(y) for x, y in kwargs.items() }
    }
  }


class OPTaskManager:


  def __init__(self, service, abandoned_poll_interval_s=ABANDONED_POLL_INTERVAL_S, data_type='op_distributed_task'):
    self._service = service
    self._data_type = data_type

    # Note that we never un-register task workers
    self._registered_task_workers = {}
    self._registered_task_workers_lock = Lock()

    self._abandoned_poll_interval_s = abandoned_poll_interval_s
    self._thread_local = threading.local()

    # TODO: you still need to create these indexes...
    # service.db.init(
    #   data_type,
    #   computed=[ ComputedIndexField('is_completed', lambda x: 'end' in x) ],
    #   indexes=[ Index('is_completed', 'deadline'), Index('parent', 'task_type_id', 'p_input') ]
    # )

    # Note: there is some risk in starting this thread right away - we might find an abandoned task that has not YET
    # been registered, but may be in the future. However, we do have up to the time specified by the
    # 'abandoned_poll_interval_s' duration, which seems pretty safe.
    service.init(lambda: Thread(target=self._poll_abandoned_tasks, daemon=True).start())


  @contextmanager
  def _exc_handler(self, msg):
    with self._service.alarm.report_exceptions(msg, log=logger, metadata=OP_DEV_METADATA):
      yield


  # Should be called from an 'except' context
  def _report_task_exception(self, msg, task_run_id, task):
    _, exc_value, exc_traceback = sys.exc_info()
    stack_summary = traceback.extract_tb(exc_traceback)
    issue_exc = { 'msg': repr(exc_value), 'traceback': stack_summary.format() }
    issue_task = _build_issue_task(task_run_id, task)
    self._service.alarm.report_issue(msg, exception=issue_exc, task=issue_task, metadata=OP_DEV_METADATA)


  def _poll_abandoned_tasks(self):
    while True:
      # TODO: interrupt this sleep when a worker becomes available if we had no capacity available last we checked
      time.sleep(self._abandoned_poll_interval_s)
      logger.debug('Checking for abandoned tasks')
      with self._exc_handler('Exception handling abandoned tasks'):
        for task_run_id, task in self._db.iter(self._data_type, Field('is_completed') == False, Field('deadline') < _now()):
          # We found an abandoned task
          if task['task_type_id'] in self._registered_task_workers:
            # We know how to complete this task
            task_worker = self._registered_task_workers[task['task_type_id']]
            with task_worker._get_capacity() as capacity:
              if capacity > 0:
                delay_s = (task_worker._heartbeat_s or MIN_HEARTBEAT_S) + HEARTBEAT_BUFFER_S
                previous_deadline = task.get('deadline')
                task['deadline'] = _now() + timedelta(seconds=delay_s)
                task['n_beats'] = 0
                try:
                  self._db.update(self._data_type, task_run_id, task, Field('is_completed') == False, Field('deadline') == previous_deadline)
                except NotFoundException:
                  # This task has been picked up by someone else
                  logger.debug('Task already picked up with ID: %s' % task_run_id)
                  continue
                logger.debug('Executing abandoned task with ID: %s' % task_run_id)
                task_worker._execute_task(task_run_id, task)
              else:
                # We may not have capacity to take on any new tasks for this worker, but we might for some other
                # worker types
                # TODO: improve this log message - we should just log once at the end of our loop
                msg = 'No capacity available to execute abandoned tasks for task type ID: %s' % task['task_type_id']
                logger.debug(msg)
                continue
          else:
            # We don't know how to complete this task
            msg = _build_abandoned_task_msg(task)
            logger.critical(msg)
            issue_task = _build_issue_task(task_run_id, task)
            self._service.alarm.report_issue(msg, task=issue_task, metadata=OP_DEV_METADATA)


  @contextmanager
  def _parent(self, task_run_id):
    self._thread_local.parent = task_run_id
    try:
      yield
    finally:
      self._thread_local.parent = None


  def _get_parent(self):
    return getattr(self._thread_local, 'parent', None)


  def __call__(self, task_type_id=None, max_workers=64, heartbeat_s=None):
    def decorator(fn):
      nonlocal task_type_id, max_workers, heartbeat_s
      if not task_type_id:
        task_type_id = _generate_task_type_id(fn)
      with self._registered_task_workers_lock:
        if task_type_id in self._registered_task_workers:
          raise Exception('Task type already registered: %s' % task_type_id)
        task_worker = TaskWorker(self, task_type_id, max_workers, heartbeat_s, fn)
        self._registered_task_workers[task_type_id] = task_worker
      return task_worker
    return decorator



# TEST
# logging.basicConfig(level=logging.DEBUG)

# from op_storage import factory
# from uuid import UUID

# config = ...
# db = factory.build('mongo_db', config)
# # db._remove_data_type('op_distributed_task')
# task = OPTaskManager(db, { 'pkg_id': None, 'config_id': None, 'commit_id': None }, abandoned_poll_interval_s=2)


# def print_tasks():
#   print('--------------------')
#   for task_id, task in db.list('op_distributed_task'):
#     print(task['task_type_id'], task_id, task['start'].isoformat(), task['end'].isoformat() if 'end' in task else None, 'parent:', task['parent'])
#   print('--------------------')
# print_tasks()


# @task(task_type_id='workflow_1', heartbeat_s=1)
# def workflow_1(x_0):
#   print('workflow_1 start')
#   x_1 = step_1(x_0 + ':1')
#   x_2 = step_2(x_1 + ':2')
#   x_3 = step_3(x_2 + ':3')
#   print('workflow_1 finish')
#   return x_3

# @task(task_type_id='step_1', heartbeat_s=1)
# def step_1(x):
#   print('step_1 start')
#   time.sleep(5)
#   print('step_1 finish')
#   return x + ':step_1'

# @task(task_type_id='step_2', heartbeat_s=1)
# def step_2(x):
#   print('step_2 start')
#   time.sleep(5)
#   print('step_2 finish')
#   return x + ':step_2'

# @task(task_type_id='step_3', heartbeat_s=1)
# def step_3(x):
#   print('step_3 start')
#   time.sleep(5)
#   print('step_3 finish')
#   return x + ':step_3'

# # workflow_1.submit('A')
# time.sleep(20)
