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, 'time.time', 'time.time', False, 'import time\n'), (106, 'os.path.isfile', 'os.path.isfile', False, 'import os\n'), (34, 'os.path.join', 'os.path.join', False, 'import os\n'), (44, 'os.path.join', 'os.path.join', False, 'import os\n'), (58, 'sample_hyperparameters.TrainableSampleGenerator', 'TrainableSampleGenerator', False, 'from sample_hyperparameters import TrainableSampleGenerator\n'), (123, 'os.path.isfile', 'os.path.isfile', False, 'import os\n'), (124, 'os.remove', 'os.remove', False, 'import os\n'), (130, 'filelock.Timeout', 'Timeout', False, 'from filelock import FileLock, Timeout\n'), (112, 'time.sleep', 'time.sleep', False, 'import time\n'), (111, 'time.time', 'time.time', False, 'import time\n')]