from contextlib import contextmanager
import time
from pathlib import Path
import random
import signal
import atexit
import sys

def release(*args):
    try:
        for f in GLOBAL_NFS_LIST.keys():
            if f.exists():
                f.unlink()

    except BaseException as exception:
        pass

def exit_release(*args):
    release()
    sys.exit(1)

atexit.register(release)
signal.signal(signal.SIGTERM, exit_release)
signal.signal(signal.SIGINT, exit_release)

GLOBAL_NFS_LIST = {}

class NFSLock:
    def __init__(self, filename, tries=float('inf')):
        self.filename = Path(filename).parent / ('.lock_' + filename.name)
        if not self.filename.exists():
            self.filename.touch()

        assert self.filename.exists()
        # generate a random 4 character string
        four_char_str = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz1234567890', k=4))
        self._start_of_name = self.filename.name + f'.attempt_'
        self._link_parent = self.filename.parent

        end_of_link = f'{time.time()}_' + four_char_str

        self.random_link = self._link_parent / f'{self._start_of_name}{end_of_link}'
        self.max_tries = tries
        GLOBAL_NFS_LIST[self.random_link] = True

    def __enter__(self):
        tries = 0
        has_alerted = False
        while True:
            self.random_link.hardlink_to(self.filename)
            num_links = self.random_link.stat().st_nlink
            if num_links == 2:
                return self.random_link

            self.random_link.unlink()

            tries += 1

            if self.max_tries is not None:
                if tries >= self.max_tries:
                    print('>> Concerning lock at', self.random_link)
                    raise TimeoutError(f'Could not acquire lock after {tries} tries..')
            elif not has_alerted:
                if tries % 5 == 0:
                    print(f'>> Failed ({tries}x) to acquire lock @', self.random_link)
                    print(f'>> Trying again!')

            time.sleep(random.random() * 5)
            # try to remove existing locks if they are older than an hour
            if tries % 10 == 0:
                parent = self._link_parent
                glob_pat = self._start_of_name + '*'
                other_locks = list(parent.glob(glob_pat))
                for other_lock in other_locks:
                    # 60 seconds per 60 min
                    if time.time() - float(other_lock.name.split('_')[-2]) > 60 * 60:
                        print('removing lock', other_lock)
                        try:
                            other_lock.unlink()
                        except:
                            print('could not remove lock', other_lock, 'arleady removed.')

    def __exit__(self, exc_type, exc_value, traceback):
        if self.random_link.exists():
            self.random_link.unlink()

        assert not self.random_link.exists()
        if self.random_link in GLOBAL_NFS_LIST:
            del GLOBAL_NFS_LIST[self.random_link]