import logging
from functools import wraps
from threading import Lock, Semaphore, Thread


DEFAULT_THREAD_POOL_SIZE = 64


# TODO: this might be a more useful interface:
# a, b = gather(Thread(target=fn_a, args=(), kwargs={}), Thread(target=fn_b, args=(), kwargs={}), max_n_threads=None)


def op_async(logger=None, thread_pool_size=DEFAULT_THREAD_POOL_SIZE):
  """
  A decorator that users the python 'Thread' object to run tasks asynchronously - note that because of the GIL, this
  decorator is only useful for non-compute tasks, like network calls. A call to the decorated function will return a
  special subclass of the thread object containing an 'wait' function. Use this function to get the results of the
  asynchronous call (or throw the exception if there is one).

  Example 1 (will complete in ~2 seconds because of the limitation on the thread_pool_size):
  @op_async(thread_pool_size=2)
  def test(i):
    import time
    time.sleep(1)
    return i
  threads = [test(i) for i in range(3)]
  results = [thread.wait() for thread in threads] # [0, 1, 2]

  Example 2 (throws an exception only when 'wait' is called):
  @op_async()
  def test():
    raise Exception()
  thread = test()
  thread.wait() # Exception is not raised until this line
  """

  if (type(thread_pool_size) is not int) or (thread_pool_size < 1):
    raise Exception("thread_pool_size must be >= 1")

  semaphore = Semaphore(thread_pool_size)
  def decorator(f):
    @wraps(f)
    def run(*args, **kwargs):
      
      class AwaitThread(Thread):

        def run(self):
          semaphore.acquire()
          try:
            self.result = f(*args, **kwargs)
          except Exception as e:
            if logger:
              logger.exception("Encountered exception in 'op_async' task")
            self.exception = e
          finally:
            semaphore.release()

        def wait(self):
          self.join()
          try:
            return self.result
          except:
            raise self.exception from None

      thread = AwaitThread()
      thread.start()
      return thread
    return run
  return decorator


def async_group(*wait_on):
  """
  Uses the 'op_async' pattern to execute the functions passed via the 'wait_on' parameter simultaneously. Returns a
  generator that delivers the result of the next sequential function as soon as it is available, or throws an exception.

  Parameters:
    'wait_on' is a list of zero-argument callables

  Example 1:
  import time
  def zero_argument_callable():
    time.sleep(1)
    return time.time()
  get = wait_group(zero_argument_callable, zero_argument_callable)
  for result in gen:
    print(result)

  Example 2 (if you don't care about the results):
  get = wait_group(zero_argument_callable, zero_argument_callable)
  tuple(gen)
  """

  threads = []
  for f in wait_on:
    thread = op_async()(f)()
    threads.append(thread)
  for thread in threads:
    yield thread.wait()


def synchronized(f):
  lock = Lock()
  def run(*args, **kwargs):
    lock.acquire()
    try:
      return f(*args, **kwargs)
    finally:
      lock.release()
  return run
