'''
Utility for simple distribution of work on multiple processes, by
making sure only one process is working on a job at once.
'''

import os
import errno
import socket
import atexit
import time
import sys


def reserve_dir(*args, exit=True):
    '''
    Convenience function to get exclusive access to a working
    experiment directory.  Exits the program if the directory is
    already done or busy (using exit_if_job_done).  Otherwise,
    returns a function creates filenames within that directory.

    Usage:

    ```
    # At the beginning of the program: exits right away if already done or being done.
    rdir = reserve_dir(experiment_directory)

    # In the middle of the program: write some files in the directory.
    with open(rdir('my-results.txt'), 'w') as f:
        f.write(mydata)

    # At the end of the program, mark the directory as done.
    rdir.done()
    ```
    '''
    directory = os.path.join(*args)
    if not exit_if_job_done(directory, exit=exit):
        return False

    def dirfn(*fn):
        return os.path.join(directory, *fn)
    dirfn.dir = directory

    def done():
        mark_job_done(directory)
    dirfn.done = done
    print('Working in %s' % directory)
    return dirfn


# Old function name.
exclusive_dirfn = reserve_dir


def mark_job_done(directory):
    '''
    Function to write a `done.txt` file into the given directory.
    Accessible as rdir.done() for rdir returned from reserve_dir.
    '''
    with open(os.path.join(directory, 'done.txt'), 'w') as f:
        f.write('done by %d@%s %s at %s' %
                (os.getpid(), socket.gethostname(),
                 os.getenv('STY', ''),
                 time.strftime('%c')))

def exit_if_job_done(directory, redo=False, force=False, verbose=True,
        exit=True):
    '''
    Convenience function to get exclusive access to an unfinished
    experiment directory.  Exits the program if the directory is
    already done or busy (using exit_of_job_done).
    '''
    if pidfile_taken(os.path.join(directory, 'lockfile.pid'),
                     force=force, verbose=verbose):
        if exit:
            sys.exit(0)
        else:
            return False
    donefile = os.path.join(directory, 'done.txt')
    if os.path.isfile(donefile):
        with open(donefile) as f:
            msg = f.read()
        if redo or force:
            if verbose:
                print('Removing %s %s' % (donefile, msg))
            os.remove(donefile)
        else:
            if verbose:
                print('%s %s' % (donefile, msg))
            if exit:
                sys.exit(0)
            else:
                return False
    return True



def pidfile_taken(path, verbose=False, force=False):
    '''
    Usage.  To grab an exclusive lock for the remaining duration of the
    current process (and exit if another process already has the lock),
    do this:

    if pidfile_taken('job_423/lockfile.pid', verbose=True):
        sys.exit(0)

    To do a batch of jobs, just run a script that does them all on
    each available machine, sharing a network filesystem.  When each
    job grabs a lock, then this will automatically distribute the
    jobs so that each one is done just once on one machine.
    '''

    # Try to create the file exclusively and write my pid into it.
    try:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_RDWR)
    except OSError as e:
        if e.errno == errno.EEXIST:
            # If we cannot because there was a race, yield the conflicter.
            conflicter = 'race'
            try:
                with open(path, 'r') as lockfile:
                    conflicter = lockfile.read().strip() or 'empty'
            except:
                pass
            # Force is for manual one-time use, for deleting stale lockfiles.
            if force:
                if verbose:
                    print('Removing %s from %s' % (path, conflicter))
                os.remove(path)
                return pidfile_taken(path, verbose=verbose, force=False)
            if verbose:
                print('%s held by %s' % (path, conflicter))
            return conflicter
        else:
            # Other problems get an exception.
            raise
    # Register to delete this file on exit.
    lockfile = os.fdopen(fd, 'r+')
    atexit.register(delete_pidfile, lockfile, path)
    # Write my pid into the open file.
    lockfile.write('%d@%s %s\n' % (os.getpid(), socket.gethostname(),
                                   os.getenv('STY', '')))
    lockfile.flush()
    os.fsync(lockfile)
    # Return 'None' to say there was not a conflict.
    return None


def delete_pidfile(lockfile, path):
    '''
    Runs at exit after pidfile_taken succeeds.
    '''
    if lockfile is not None:
        try:
            lockfile.close()
        except:
            pass
    try:
        os.unlink(path)
    except:
        pass
