from __future__ import print_function

import six

import copy
from functools import partial
import json
import math
import random
import signal
import threading
import time
import traceback


import gevent.pool
import multiprocessing as mp
import numpy as np
import pathos.pools
from tqdm import tqdm

from .threaded_queue import ThreadedQueue
from .lib import yellow, scatter, sip, toiter

from .aws_queue_api import AWSTaskQueueAPI, AWS_BATCH_SIZE
from .file_queue_api import FileQueueAPI
from .paths import extract_path, mkpath
from .registered_task import RegisteredTask, deserialize, totask, totaskid
from .scheduler import schedule_jobs
from .secrets import (
  AWS_DEFAULT_REGION
)

def totalfn(iterator, total):
  if total is not None:
    return total
  try:
    return len(iterator)
  except TypeError:
    return None

class QueueEmptyError(LookupError):
  pass

LEASE_SECONDS = 300

class TaskQueue(object):
  """
  The standard usage is that a client calls lease to get the next available task,
  performs that task, and then calls task.delete on that task before the lease expires.
  If the client cannot finish the task before the lease expires,
  and has a reasonable chance of completing the task,
  it should call task.update before the lease expires.
  If the client completes the task after the lease has expired,
  it still needs to delete the task. 
  Tasks should be designed to be idempotent to avoid errors 
  if multiple clients complete the same task.
  """
  def __init__(
    self, qurl, n_threads=40, 
    green=False, progress=True, 
    **kwargs
  ):
    self.qurl = qurl
    self.path = extract_path(qurl)
    self.api = self.initialize_api(self.path, kwargs)
    self.n_threads = n_threads
    self.green = bool(green)
    self.progress = bool(progress),
    self.kwargs = kwargs

    if self.green:
      self.check_monkey_patch_status()

  @property
  def qualified_path(self):
    return mkpath(self.path)

  def initialize_api(self, path, kwargs):
    if path.protocol == 'sqs':
      return AWSTaskQueueAPI(qurl=path.path, region_name=kwargs.get('region', AWS_DEFAULT_REGION))
    elif path.protocol == 'fq':
      return FileQueueAPI(path.path)
    else:
      raise ValueError('Unsupported protocol ' + str(self.path.protocol))

  def check_monkey_patch_status(self):
    import gevent.monkey
    if not gevent.monkey.is_module_patched("socket"):
      print(yellow("""
    Green threads require monkey patching the standard library 
    to use a non-blocking network socket call.

    Please place the following lines at the beginning of your
    program. `thread=False` is there because sometimes this
    causes hanging in multiprocessing.

    import gevent.monkey
    gevent.monkey.patch_all(thread=False)
        """))

  @property
  def enqueued(self):
    """
    Returns the approximate(!) number of tasks enqueued in the cloud.
    WARNING: The number computed by Google is eventually
      consistent. It may return impossible numbers that
      are small deviations from the number in the queue.
      For instance, we've seen 1005 enqueued after 1000 
      inserts.
    
    Returns: (int) number of tasks in cloud queue
    """
    return self.api.enqueued

  @property
  def inserted(self):
    return self.api.inserted

  @property
  def completed(self):
    return self.api.completed

  def is_empty(self):
    return self.api.is_empty()

  # def status(self):
  #   """
  #   Gets information about the TaskQueue
  #   """
  #   return self.api.get(getStats=True)

  def list(self):
    """
    Lists all non-deleted Tasks in a TaskQueue, 
    whether or not they are currently leased, 
    up to a maximum of 100.
    """
    return [ totask(x) for x in iter(self.api) ]

  def insert(self, tasks, delay_seconds=0, total=None, parallel=1):
    tasks = toiter(tasks)
    total = totalfn(tasks, total)

    if parallel not in (1, False) and total is not None and total > 1:
      multiprocess_upload(self.__class__, mkpath(self.path), tasks, parallel=parallel)
      return

    try:
      batch_size = self.api.batch_size
    except:
      batch_size = 1
    
    bodies = (
      {
        "payload": task.payload(),
        "queueName": self.path.path,
      } 
      for task in tasks
    )

    ct = [ 0 ] # using a list instead of a raw number to use pass-by-reference
    ct_lock = threading.Lock()
    def insertfn(batch, ct):
      try:
        incr = self.api.insert(batch, delay_seconds) 
      finally:
        with ct_lock:
          ct[0] += incr
    
    schedule_jobs(
      fns=( partial(insertfn, batch, ct) for batch in sip(bodies, batch_size) ),
      concurrency=self.n_threads,
      progress='Inserting',
      total=total,
      green=self.green,
      batch_size=batch_size,
    )

    self.api.add_insert_count(ct[0])

  def insert_all(self, *args, **kwargs):
    """For backwards compatibility."""
    return self.insert(*args, **kwargs)

  def rezero(self):
    """Resets statistic counters such as completions and insertions to zero."""
    self.api.rezero()

  def renew(self, task, seconds):
    """Update the duration of a task lease."""
    return self.api.renew_lease(task, seconds)

  def cancel(self, task):
    return self.api.cancel_lease(task)

  def release_all(self):
    return self.api.release_all()

  def lease(self, seconds=600, num_tasks=1):
    """
    Acquires a lease on the topmost N unowned tasks in the specified queue.
    Required query parameters: leaseSecs, numTasks
    """
    if num_tasks <= 0:
      raise ValueError("num_tasks must be > 0. Got: " + str(num_tasks))
    if seconds < 0:
      raise ValueError("lease seconds must be >= 0. Got: " + str(seconds))

    tasks = self.api.lease(seconds, num_tasks)

    if not len(tasks):
      raise QueueEmptyError()

    if num_tasks == 1:
      return totask(tasks[0])
    else:
      return [ totask(task) for task in tasks ]

  def delete(self, task_id, total=None, tally=False):
    """Deletes a task from a TaskQueue."""
    task_id = toiter(task_id)
    total = totalfn(task_id, total)

    def deltask(tid):
      self.api.delete(totaskid(tid))
      if tally:
        self.api.tally()

    schedule_jobs(
      fns=( partial(deltask, tid) for tid in task_id ),
      concurrency=self.n_threads,
      progress=('Deleting' if self.progress else None),
      total=total,
      green=self.green,
    )

  def purge(self):
    """Deletes all tasks in the queue."""
    try:
      return self.api.purge()
    except AttributeError:
      while True:
        lst = self.list()
        if len(lst) == 0:
          break

        for task in lst:
          self.delete(task)
        self.wait()
      return self

  def poll(
    self, lease_seconds=LEASE_SECONDS,  
    verbose=False, execute_args=[], execute_kwargs={}, 
    stop_fn=None, backoff_exceptions=[], min_backoff_window=1, 
    max_backoff_window=120, before_fn=None, after_fn=None,
    tally=False
  ):
    """
    Poll a queue until a stop condition is reached (default forever). Note
    that this function is not thread safe as it requires a global variable
    to intercept SIGINT.
    lease_seconds: each task should be leased for this many seconds
    execute_args / execute_kwargs: pass these arguments to task execution
    backoff_exceptions: A list of exceptions that instead of causing a crash,
      instead cause the polling to back off for an increasing exponential 
      random window.
    min_backoff_window: The minimum sized window (in seconds) to select a 
      random backoff time. 
    max_backoff_window: The window doubles each retry. This is the maximum value
      in seconds.
    stop_fn: A boolean returning function that accepts no parameters. When
      it returns True, the task execution loop will terminate. It is evaluated
      once after every task. If you provide the arguments `executed` (tasks completed) 
      `tries` (current attempts at fetching a task), or `previous_execution_time` (time in
      seconds to run the last task) they will be dependency injected.
    before_fn: Pass task pre-execution.
    after_fn: Pass task post-execution.
    verbose: print out the status of each step
    tally: contribute each completed task to a completions counter if supported.

    Return: number of tasks executed
    """
    global LOOP

    if not callable(stop_fn) and stop_fn is not None:
      raise ValueError("stop_fn must be a callable. " + str(stop_fn))
    elif not callable(stop_fn):
      stop_fn = lambda: False

    def random_exponential_window_backoff(n):
      n = min(n, min_backoff_window)
      # 120 sec max b/c on avg a request every ~250msec if 500 containers 
      # in contention which seems like a quite reasonable volume of traffic 
      # to handle
      high = min(2 ** n, max_backoff_window) 
      return random.uniform(0, high)

    def printv(*args, **kwargs):
      if verbose:
        print(*args, **kwargs)

    LOOP = True

    def sigint_handler(signum, frame):
      global LOOP
      printv("Interrupted. Exiting after this task completes...")
      LOOP = False
    
    prev_sigint_handler = signal.getsignal(signal.SIGINT)
    signal.signal(signal.SIGINT, sigint_handler)

    tries = 0
    executed = 0

    backoff = False
    backoff_exceptions = tuple(list(backoff_exceptions) + [ QueueEmptyError ])

    before_fn = before_fn or (lambda x: x)
    after_fn = after_fn or (lambda x: x)

    while LOOP:
      task = 'unknown' # for error message prior to leasing
      try:
        task = self.lease(seconds=int(lease_seconds))
        tries += 1
        before_fn(task)
        time_start = time.time()
        task.execute(*execute_args, **execute_kwargs)
        time_delta = time.time() - time_start
        executed += 1
        printv("Delete enqueued task...")
        self.delete(task, tally=tally)
        printv('INFO', task , "succesfully executed in {:.2f} sec.".format(time_delta))
        after_fn(task)
        tries = 0
      except backoff_exceptions:
        backoff = True
      except Exception as e:
        printv('ERROR', task, "raised {}\n {}".format(e , traceback.format_exc()))
        raise # this will restart the container in kubernetes
      
      varnames = stop_fn.__code__.co_varnames
      stop_fn_bound = stop_fn
      if 'executed' in varnames:
        stop_fn_bound = partial(stop_fn_bound, executed=executed)
      if 'tries' in varnames:
        stop_fn_bound = partial(stop_fn_bound, tries=tries)
      if 'previous_execution_time' in varnames:
        stop_fn_bound = partial(stop_fn_bound, previous_execution_time=time_delta) 

      if stop_fn_bound():
        break

      if backoff:
        time.sleep(random_exponential_window_backoff(tries))

      backoff = False

    printv("Task execution loop exited.")
    signal.signal(signal.SIGINT, prev_sigint_handler)

    return executed

  def block_until_empty(self, interval_sec=2):
    while self.enqueued > 0:
      time.sleep(interval_sec)

  def __enter__(self):
    return self

  def __exit__(self, exception_type, exception_value, traceback):
    pass

class LocalTaskQueue(object):
  def __init__(self, parallel=1, queue_name='', queue_server='', progress=True):
    if parallel and type(parallel) == bool:
      parallel = mp.cpu_count()

    self.parallel = parallel
    self.queue = []
    self.progress = progress

  def insert(
      self, tasks, 
      delay_seconds=0, total=None,
      parallel=None, progress=True
    ):
    for task in tasks:
      args, kwargs = [], {}
      if isinstance(task, tuple):
        task, args, kwargs = task
      task = {
        'payload': task.payload(),
        'id': -1,
      }
      self.queue.append( (task, args, kwargs) )

  def insert_all(self, *args, **kwargs):
    self.insert(*args, **kwargs)
    self.execute(self.progress)

  def poll(self, *args, **kwargs):
    pass

  def execute(self, progress=True, parallel=None, total=None):
    if parallel is None:
      parallel = self.parallel

    total = totalfn(self.queue, total)

    with tqdm(total=total, desc="Tasks", disable=(not progress)) as pbar:
      if self.parallel == 1:
        while self.queue:
          _task_execute(self.queue.pop(0))
          pbar.update()
      else:
        with pathos.pools.ProcessPool(self.parallel) as executor:
          for _ in executor.imap(_task_execute, self.queue):
            pbar.update()
    
      self.queue = []

  def __enter__(self):
    return self

  def __exit__(self, exception_type, exception_value, traceback):
    self.execute()

class MockTaskQueue(LocalTaskQueue):
  pass

class GreenTaskQueue(TaskQueue):
  def __init__(self, *args, **kwargs):
    kwargs['green'] = True
    super(GreenTaskQueue, self).__init__(*args, **kwargs)

# Necessary to define here to make the 
# function picklable
def _task_execute(task_tuple):
  task, args, kwargs = task_tuple
  task = totask(task)
  task.execute(*args, **kwargs)

## Multiprocess Upload

def _scatter(sequence, n):
  """Scatters elements of ``sequence`` into ``n`` blocks."""
  chunklen = int(math.ceil(float(len(sequence)) / float(n)))

  return [
    sequence[ i*chunklen : (i+1)*chunklen ] for i in range(n)
  ]

def soloprocess_upload(QueueClass, queue_name, tasks):
  tq = QueueClass(queue_name)
  tq.insert(tasks)

error_queue = mp.Queue()

def multiprocess_upload(QueueClass, queue_name, tasks, parallel=True):
  if parallel is True:
    parallel = mp.cpu_count()
  elif parallel <= 0:
    raise ValueError("Parallel must be a positive number or zero (all cpus). Got: " + str(parallel))

  if parallel == 1:
    soloprocess_upload(QueueClass, queue_name, tasks)
    return 

  def capturing_soloprocess_upload(*args, **kwargs):
    try:
      soloprocess_upload(*args, **kwargs)
    except Exception as err:
      print(err)
      error_queue.put(err)

  uploadfn = partial(
    capturing_soloprocess_upload, QueueClass, queue_name
  )
  tasks = _scatter(tasks, parallel)

  # This is a hack to get dill to pickle dynamically
  # generated classes. This is an important use case
  # for when we create iterators with generator __iter__
  # functions on demand.

  # https://github.com/uqfoundation/dill/issues/56

  try:
    task = next(item for item in tasks if item is not None)
  except StopIteration:
    return 

  cls_module = task.__class__.__module__
  task.__class__.__module__ = '__main__'

  with pathos.pools.ProcessPool(parallel) as pool:
    pool.map(uploadfn, tasks)

  task.__class__.__module__ = cls_module

  if not error_queue.empty():
    errors = []
    while not error_queue.empty():
      err = error_queue.get()
      if err is not StopIteration:
        errors.append(err)
    if len(errors):
      raise Exception(errors)