from filelock import FileLock, Timeout
import os
import time

class ProcessFileLock(FileLock):
    """
    FileLock that is unique per path in each process (for, eg., reentrance)
    """
    locks = {}
    def __new__(cls, path, *args, **kwargs):
        if path in ProcessFileLock.locks:
            return ProcessFileLock.locks[path]
        else:
            lock = super().__new__(cls, path, *args, **kwargs)
            lock.__new_init__(path, *args, **kwargs)
            ProcessFileLock.locks[path] = lock
            return lock
            
    def __new_init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    
    def __init__(self, *args, **kwargs):
        pass

class ExplosiveFileLock(ProcessFileLock):
    def acquire(self, *args, **kwargs):
        r = super().acquire(*args, **kwargs)
        if self._lock_counter > 1:
            raise BlockingIOError(f"Process attempted to reacquire lock for {self._lock_file}")
        return r
       
class HistoriesLock(FileLock):
    def __init__(self, dir, ensemble=None):
        super().__init__(os.path.join(dir, "histories.lock"))
        self.ensemble = ensemble

    def release(self, **kwargs):
        super().release()
        if self.ensemble and self._lock_counter == 0:
            self.ensemble.close_histories()

class SamplesLock(FileLock):    
    def __init__(self, dir, ensemble=None):
        super().__init__(os.path.join(dir, "samples.lock"))
        self.ensemble = ensemble

    def release(self, **kwargs):
        if self.ensemble and self._lock_counter == 1:
            self.ensemble._test_samples.close()
            self.ensemble._test_samples = None
        super().release()
            
    def __enter__(self):
        print("Acquiring samples lock... ", end='')
        super().__enter__()
        if self.ensemble._test_samples is None:
            from sample_hyperparameters import TrainableSampleGenerator
            self.ensemble._test_samples = TrainableSampleGenerator(self.ensemble.config.exploration.sample_file, configs=self.ensemble.config_files, stub=self.ensemble.stub)
        print("Done.")
        return self.ensemble._test_samples

    
class ExistLock:
    """
    Locks on the existence of the given file.
    No guarantees of atomicity!
    Unique per process, for reentry
    """
    locks={}
    def __new__(cls, path, *args, **kwargs):
        if path in ExistLock.locks:
            lock = ExistLock.locks[path]
            #print(f"Reloading ExistLock('{path}')")
            #print(f"    Lock counter = {lock._lock_counter}")
            return lock
        else:
            #print(f"Creating new ExistLock('{path}')")
            lock = super().__new__(cls)
            lock.__new_init__(path, *args, **kwargs)
            ExistLock.locks[path] = lock
            return lock
    
    def __new_init__(self, path, block=True, timeout=None, polling_interval=.05):
        self.path = path
        if not block:
            timeout == 0.0
        else:
            self.timeout=timeout
        self.polling_interval=polling_interval
        self._lock_counter = 0
    
    def acquire(self, block=None, timeout=None):
        """
        Not atomic. Should probably happen within the context of an
        atomic lock.
        """
        if block == False:
            timeout = 0.0
        if timeout is None:
            timeout = self.timeout
            
        #print(f"Trying to acquire ExistLock('{self.path}')...")
        #print(f"    Lock counter = {self._lock_counter}")

        start_time = time.time()
        while os.path.isfile(self.path):
            if self._lock_counter > 0:
                self._lock_counter += 1
                #print(f"Acquired, lock counter = {self._lock_counter}")
                return True
            if timeout is None or time.time() - start_time < timeout:
                time.sleep(self.polling_interval)
            else:
                return False
        
        with open(self.path, 'w'):
            self._lock_counter = 1
            #print(f"Acquired, lock counter = {self._lock_counter}")
            return True
    
    def release(self):
        self._lock_counter = min(0, self._lock_counter - 1)
        if self._lock_counter == 0 and os.path.isfile(self.path):
            os.remove(self.path)
            
    def __enter__(self):
        if self.acquire():
            return self
        else:
            raise Timeout(f"Failed to acquire ExistLock for file {self.path}")
            
    def __exit__(self, type, value, traceback):
        self.release()


[((105, 21, 105, 32), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((106, 14, 106, 39), 'os.path.isfile', 'os.path.isfile', ({(106, 29, 106, 38): 'self.path'}, {}), '(self.path)', False, 'import os\n'), ((34, 25, 34, 60), 'os.path.join', 'os.path.join', ({(34, 38, 34, 41): 'dir', (34, 43, 34, 59): '"""histories.lock"""'}, {}), "(dir, 'histories.lock')", False, 'import os\n'), ((44, 25, 44, 58), 'os.path.join', 'os.path.join', ({(44, 38, 44, 41): 'dir', (44, 43, 44, 57): '"""samples.lock"""'}, {}), "(dir, 'samples.lock')", False, 'import os\n'), ((58, 42, 58, 173), 'sample_hyperparameters.TrainableSampleGenerator', 'TrainableSampleGenerator', (), '', False, 'from sample_hyperparameters import TrainableSampleGenerator\n'), ((123, 39, 123, 64), 'os.path.isfile', 'os.path.isfile', ({(123, 54, 123, 63): 'self.path'}, {}), '(self.path)', False, 'import os\n'), ((124, 12, 124, 32), 'os.remove', 'os.remove', ({(124, 22, 124, 31): 'self.path'}, {}), '(self.path)', False, 'import os\n'), ((130, 18, 130, 78), 'filelock.Timeout', 'Timeout', ({(130, 26, 130, 77): 'f"""Failed to acquire ExistLock for file {self.path}"""'}, {}), "(f'Failed to acquire ExistLock for file {self.path}')", False, 'from filelock import FileLock, Timeout\n'), ((112, 16, 112, 49), 'time.sleep', 'time.sleep', ({(112, 27, 112, 48): 'self.polling_interval'}, {}), '(self.polling_interval)', False, 'import time\n'), ((111, 34, 111, 45), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n')]