import os
import shelve

from contextlib import suppress
from fasteners import InterProcessLock
from threading import RLock, ThreadError

from .single_process import DistributedTaskManagerSingleProcess


class _LockContextManager(object):

  def __init__(self, path, task_manager):
    self.task_manager = task_manager
    os.makedirs(path, exist_ok=True)
    self.thread_lock = RLock()
    lock_file_path = os.path.join(path, "lock")
    self.process_lock = InterProcessLock(lock_file_path)
    self.data_file_path = os.path.join(path, "data")
    self.ref_count = 0

  def __enter__(self):
    self.thread_lock.acquire()
    self.process_lock.acquire()
    if self.ref_count == 0:
      self.task_manager.tasks = shelve.open(self.data_file_path)
    self.ref_count += 1

  def __exit__(self, *exc):
    self.ref_count -= 1
    if self.ref_count == 0:
      self.task_manager.tasks.close()
    with suppress(ThreadError):
      # Suppressing errors here effectively makes this a re-entrant lock
      self.process_lock.release()
    self.thread_lock.release()


class DistributedTaskManagerShelf(DistributedTaskManagerSingleProcess):

  def __init__(self, path):
    super().__init__()
    self.lock = _LockContextManager(path, self)
