#! /usr/bin/env python

# This utility helps with copying select files from one directory to another.
# You provide it with a source directory and a destination directory, and it
# shows a numbered listing of each file in the source directory, from which
# list you can choose which files to copy to the destination directory. It does
# this recursively through the directory structure, and when you've told it all
# the files you want, it copies them to the new location.
#
# The standard mode is an interactive mode with various commands for expediting
# the task, such as quickly navigating through the directory structure,
# appending files to the queue, listing and clearing the queue, etc. A list of
# these commands is available by typing h at the prompt.
#
# By default, the original directory structure under srcdir is maintained as
# files are copied over to destdir. You can override this behavior with the
# --no-subdirs option, in which case all files will go directly into destdir
# with no subdirectories at all.
#
# The --one-by-one option puts you into one-by-one mode, which means you won't
# get a numbered listing of all the files as you go into each directory, but
# rather a prompt for a yes or no on each file.
#
# The -v or --verbose option will increase the amount of descriptive output you
# get. Right now this mainly just means getting a list of files that are being
# copied at the time of copying.
#
# The --queuefile option allows you to simply copy files as listed in the file
# specified with the option, and not go into normal interactive mode. The file
# can contain an optional first line of the format "srcdir: /path/to/dir". If
# it does, you can leave out the --srcdir option on the command line. If you
# choose to also include the --srcdir option, that will override the srcdir
# line in the file. The rest of the file is a list of files to be copied, one
# per line, with full paths relative to the srcdir. You can create a queuefile
# at any time in interactive mode by using the 'w' command, and then you can
# later use the queuefile with the --queuefile option as described above.
#
# This script was inspired by the need to go through a huge directory structure
# filled with mp3s and select my favorite ones for moving to another directory
# and making a cd therefrom. Sure, I could do it all by hand, but that just
# seemed like an unsavorily arduous task.
#
# With that task in mind, the queuefile functionality provides for saving lists
# of the files on a given cd so you can re-create the cd later with just a few
# commands.
#
# This script is distributed under the terms of the Gnu General Public License
# and comes with no warranty whatsoever. Please see www.gnu.org for more
# information.


import os, os.path, getopt, sys, shutil, re

#------------------------------------------------------------------------------#
class WorkingDir:
    '''Represents the current working directory. Includes properites for
    subdirectories and files in the current directory, each a list, as well as
    one for the parent directory, which will only go as high as the srcdir
    passed to the program, and one for the directory to go to if the program
    needs to advance to the next in the list.'''

    def __init__(self, dir):
        self.dir = dir

        # determine the parent directory
        if dir[-1] == '/':
            dir = dir[0:len(dir)-1]
        if dir != srcdir:
            self.parent = dir[0:dir.rfind('/')+1]
        else:
            self.parent = srcdir

        diritems = os.listdir(dir)
        diritems.sort()
        diritems = [dir + '/' + item for item in diritems]
        indices = range(1, len(diritems) + 1)

        # get subdirs from diritems
        self.subdirs = [item for item in diritems if os.path.isdir(item)]
        self.subdirs = zip(indices, self.subdirs)

        # get files and sizes from diritems
        self.files = [item for item in diritems if os.path.isfile(item)]
        filesizes = []
        for file in self.files:
            filesizes.append(int(os.path.getsize(file)))
        self.files = zip(indices, self.files, filesizes)

        # note: oboidx is a zero-based index of the files array itself - not to
        # be confused with the one-based index created by zipping indices with
        # the files list and stored as the first half of each tuple in files
        self.oboidx = 0

        # determine the next directory
        nextidx = None
        indices = xrange(len(dirlist))
        for idx, ldir in zip(indices, dirlist):
            if ldir == self.dir:
                nextidx = int(idx) + 1
                break
        if nextidx is None:
            print "couldn't find the current directory in the dirlist"
        elif nextidx < len(dirlist):
            self.next = dirlist[nextidx]
        else:
            self.next = None
#------------------------------------------------------------------------------#
def init():
    global srcdir, destdir, keep_subdirs, one_by_one, queue
    global wkdir, filesre, cdre, wre, rre, ure, dirlist, queuefile, verbose
    global interactive, obo_shown, destre, vre, sdre

    try:
        shortopts = "hd:s:v"
        longopts = ["help", "verbose", "queuefile=", "no-subdirs", "one-by-one"]
        opts, args = getopt.getopt(sys.argv[1:], shortopts, longopts)
    except getopt.GetoptError:
        show_usage()

    srcdir = None
    destdir = None
    verbose = False
    interactive = True
    queuefile = None
    keep_subdirs = True
    one_by_one = False
    obo_shown = False

    for o, a in opts:
        if o in ('-h', '--help'):
            show_usage()
        elif o in ('-s', '--srcdir'):
            srcdir = a
        elif o in ('-d', '--destdir'):
            destdir = a
        elif o in ('-v', '--verbose'):
            verbose = True
        elif o == '--queuefile':
            queuefile = a
        elif o == '--no-subdirs':
            keep_subdirs = False
        elif o == '--one-by-one':
            one_by_one = True

    if (srcdir is None and queuefile is None) or destdir is None:
        show_usage()

    if queuefile is not None:
        interactive = False
        if srcdir is None:
            srcdir = get_queuefile_srcdir()

    if not os.path.isdir(srcdir):
        sys.exit('%s is not a directory.' % (srcdir))
        
    chk_destdir()

    if srcdir[-1] != '/': srcdir += '/'
    if destdir[-1] != '/': destdir += '/'

    filesre = re.compile('^[\s\d]+$')
    cdre = re.compile('^\s*d\s*(\d+)\s*$')
    wre = re.compile('^\s*w\s+([\w/\.]+)\s*$')
    rre = re.compile('^\s*r\s+([\w/\.]+)\s*$')
    ure = re.compile('^\s*uq(\s+((\d+)|(q\d+)))+\s*$')
    destre = re.compile('^\s*set\s+destdir\s+([\w/\.]+)\s*$')
    vre = re.compile('^\s*set\s+verbose\s+(on|off)\s*$')
    sdre = re.compile('^\s*set\s+no-subdirs\s+(on|off)\s*$')

    dirlist = []
    if queuefile is None:
        extend_dirlist(srcdir)

    queue = []
#------------------------------------------------------------------------------#
def main():
    global obo_shown

    if queuefile is not None:
        read_queuefile(queuefile)
        copy_queue()
        sys.exit()

    chdir(srcdir)

    while 1:
        prompt = 'command (h for help): '

        if one_by_one:
            prompt = 'y/n or command: '
            if obo_shown is False:
                show_obo()
        else:
            print

        obo_shown = False

        cmd = raw_input(prompt).strip()

        if len(cmd) == 0:
            pass
        elif filesre.match(cmd):
            addfiles(cmd)
        elif cmd == 'a':
            add_all_files()
        elif cmd == 'l':
            showdir()
        elif cmd == 'lq':
            list_queue()
        elif cmd == 'qs' or cmd == 'qsb':
            show_queue_size(cmd)
        elif cmd == 'd':
            into_nextdir()
        elif cdre.match(cmd):
            into_subdir(cmd)
        elif cmd == 'u':
            updir()
        elif cmd == 'pwd':
            print_wkdir()
        elif cmd == 'srcdir':
            print_srcdir()
        elif cmd == 'destdir':
            print_destdir()
        elif one_by_one and ((cmd == 'y') or (cmd == 'n')):
            do_obo(cmd)
        elif cmd == 'copy':
            copy_queue()
        elif wre.match(cmd):
            write_queuefile(cmd)
        elif ure.match(cmd):
            unqueue(cmd)
        elif cmd == 'xq':
            clear_queue()
        elif rre.match(cmd):
            read_queuefile(cmd)
        elif destre.match(cmd):
            set_destdir(cmd)
        elif vre.match(cmd):
            set_verbose(cmd)
        elif sdre.match(cmd):
            set_no_subdirs(cmd)
        elif cmd == 'verbose':
            show_verbose_status()
        elif cmd == 'no-subdirs':
            show_no_subdirs_status()
        elif cmd in ['quit', 'q']:
            sys.exit('')
        elif cmd == 'h':
            show_cmd_help()
        else:
            print "didn't understand that command"
#------------------------------------------------------------------------------#
def get_queuefile_srcdir():
    cf = open(queuefile)
    sdline = cf.readline()
    cf.close()

    sdre = re.compile('^srcdir:\s*([\w/\.]+)\s*\n$')

    m = sdre.match(sdline)
    if m is not None:
        return m.group(1)
    else:
        errmsg = "srcdir line in %s is missing or incorrect" % (queuefile)
        sys.exit(errmsg)
#------------------------------------------------------------------------------#
def addfiles(cmd):
    # they specified files to add
    newfiles = [(f,s) for i,f,s in wkdir.files if str(i) in cmd.split()]

    msg = "adding file" + ('s', '')[len(newfiles) == 1] + " to queue"
    print msg

    extend_queue(newfiles)

    show_queue_size()
#------------------------------------------------------------------------------#
def add_all_files():
    # add all files in the current dir
    newfiles = [(f,s) for i,f,s in wkdir.files]

    msg = "adding file" + ('s', '')[len(newfiles) == 1] + " to queue"
    print msg

    extend_queue(newfiles)

    show_queue_size()
#------------------------------------------------------------------------------#
def extend_queue(newfiles):
    qfiles = [f for f,s in queue]
    queue.extend([(f,s) for f,s in newfiles if f not in qfiles])
#------------------------------------------------------------------------------#
def list_queue():
    if len(queue) == 0:
        print "no files queued for copying"
    else:
        print
        print 'files queued for copying:'
        print '-' * 80

        indices = range(1, len(queue)+1)
        files, sizes = [f for f,s in queue], [s for f,s in queue]
        indqueue = zip(indices, files, sizes)

        for idx, file, size in indqueue:
            print "q%s - %s - %s" % (idx, human_readable(size), file)

        print '-' * 80
        show_queue_size()
#------------------------------------------------------------------------------#
def show_queue_size(cmd='qs'):
    print 'total queue size is %s' % (get_queue_size(cmd))
#------------------------------------------------------------------------------#
def get_queue_size(cmd='qs'):
    qsize = 0
    for file, size in queue:
        qsize += size
    
    if cmd == 'qsb':
        return repr(qsize)
    else:
        return human_readable(qsize)
#------------------------------------------------------------------------------#
def human_readable(bytes):
    if bytes < 1024:
        return repr(bytes)
    elif bytes < 1048576:
        x = float(bytes / 1024.0)
        return '%.1fk' % (x)
    elif bytes < 1073741824:
        x = float(bytes / 1048576.0)
        return '%.1fM' % (x)
    else:
        x = float(bytes / 1073741824.0)
        return '%.1fG' % (x)
#------------------------------------------------------------------------------#
def into_nextdir():
    if wkdir.next is None:
        print "already in the last directory"
    else:
        chdir(wkdir.next)
#------------------------------------------------------------------------------#
def into_subdir(cmd):
    m = cdre.match(cmd)
    subdir_idx = int(m.group(1)) - 1

    if subdir_idx < len(wkdir.subdirs):
        chdir(wkdir.subdirs[subdir_idx][1])
    else:
        print "no such directory"
#------------------------------------------------------------------------------#
def updir():
    if wkdir.dir == srcdir:
        print "already at the highest directory"
    else:
        chdir(wkdir.parent)
#------------------------------------------------------------------------------#
def chdir(dir):
    global wkdir
    wkdir = WorkingDir(dir)
    showdir()
#------------------------------------------------------------------------------#
def showdir():
    print
    print 'directory: %s' % (wkdir.dir)
    print '-' * 80

    # display the numbered subdirs
    if len(wkdir.subdirs) > 0:
        print "subdirectories"
        for idx, subdir in wkdir.subdirs:
            subdir = os.path.basename(subdir)
            print 'd%s - %s/' % (str(idx), subdir)

    # display the numbered files
    if len(wkdir.files) > 0:
        if one_by_one:
            print "\nfiles (one-by-one mode)"
        else:
            print "files"
        if one_by_one:
            show_obo()
        else:
            for idx, file, size in wkdir.files:
                msg = ('%s', '[%s]')[(file, size) in queue]
                msg += ' - %s - %s'
                msg = msg % (idx, human_readable(size), os.path.basename(file))
                print msg
            print '-' * 80
#------------------------------------------------------------------------------#
def print_wkdir():
    print wkdir.dir
#------------------------------------------------------------------------------#
def print_srcdir():
    print srcdir
#------------------------------------------------------------------------------#
def print_destdir():
    print destdir
#------------------------------------------------------------------------------#
def show_obo():
    global wkdir, obo_shown

    if wkdir.oboidx < len(wkdir.files):
        idx, file, size = wkdir.files[wkdir.oboidx]
        print
        print '%s - %s' % (idx, file)
        obo_shown = True
    else:
        chdir(wkdir.next)
#------------------------------------------------------------------------------#
def do_obo(cmd):
    global wkdir

    if cmd == 'y':
        idx, file, size = wkdir.files[wkdir.oboidx]
        if file not in [f for f,s in queue]:
            queue.append((file, size))
        print "adding file"
        show_queue_size()

    wkdir.oboidx += 1

    show_obo()
#------------------------------------------------------------------------------#
def copy_queue():
    if len(queue) == 0:
        print 'no files to copy'
        return

    if verbose:
        print
        print 'COPYING:'
        print '-' * 80
    else:
        print "copying files"

    indices = range(0, len(queue))
    for idx, file in zip(indices, [f for f,s in queue]):
        if not os.path.isfile(file):
            if verbose:
                print file + " is not a file. skipping."
                if idx != len(queue) - 1:
                    print
            continue

        filename = os.path.basename(file)
        newpath = ''
        if keep_subdirs:
            path = os.path.dirname(file)
            newpath = path[len(srcdir):]
            if len(newpath) > 0:
                newpath += '/'
        newfile = destdir + newpath + filename

        if verbose:
            print 'from: %s' % (file)
            print '  to: %s' % (newfile)
            if idx != len(queue) - 1:
                print

        try:
            os.makedirs(os.path.dirname(newfile), 0755)
        except:
            pass

        shutil.copy(file, newfile)

    if verbose:
        print '-' * 80
#------------------------------------------------------------------------------#
def clear_queue():
    global queue
    print "clearing queue"
    queue = []
#------------------------------------------------------------------------------#
def unqueue(cmd):
    global queue

    cmd = cmd[2:].strip()
    filenums = cmd.split()

    uq_idxs = []

    for num in filenums:
        if num[0] == 'q':
            qidx = int(num[1:]) - 1
            uq_idxs.append(qidx)
        else:
            idx, file, size = wkdir.files[int(num)-1]
            qidx = [f for f,s in queue].index(file)
            uq_idxs.append(qidx)

    print "un-queueing specified files"
    indqueue = zip(range(0, len(queue)), queue)
    queue = [f_s for i,f_s in indqueue if i not in uq_idxs]
#------------------------------------------------------------------------------#
def write_queuefile(cmd):
    m = wre.match(cmd)
    filename = m.group(1)
    
    try:
        cf = open(filename, 'w')
    except IOError:
        print "can't open %s for writing" % (filename)
        return

    print "writing queuefile"

    cf.write("srcdir: %s\n\n" % (srcdir))

    for file, size in queue:
        file = file[len(srcdir):]
        cf.write(file + "\n")

    cf.close()
#------------------------------------------------------------------------------#
def read_queuefile(inp):
    if rre.match(inp):
        m = rre.match(inp)
        filename = m.group(1)
    else:
        filename = inp
    
    try:
        qf = open(filename, 'r')
    except IOError:
        print "can't open %s for reading" % (filename)
        return

    if interactive is True:
        print "appending files from queuefile"

    qf_lines = qf.readlines()
    qfiles = [f for f,s in queue]
    newfiles = []

    for line in qf_lines:
        if line[-1] == '\n':
            line = line[0:len(line)-1]
        if line[0:7] == 'srcdir:' or len(line.strip()) == 0:
            continue
        file = srcdir + line
        if file not in qfiles:
            if os.path.isfile(file):
                newfiles.append((file, int(os.path.getsize(file))))

    qf.close()

    extend_queue(newfiles)
    
    show_queue_size()
#------------------------------------------------------------------------------#
def extend_dirlist(dir):
    if dir[-1] != '/':
        dir += '/'

    dirlist.append(dir)

    diritems = os.listdir(dir)
    diritems.sort()
    diritems = [dir + item for item in diritems]

    # get the subdirs
    subdirs = [item for item in diritems if os.path.isdir(item)]

    for subdir in subdirs:
        extend_dirlist(subdir)
#------------------------------------------------------------------------------#
def set_destdir(cmd):
    global destdir

    m = destre.match(cmd)
    newdir = m.group(1)

    gooddir = chk_destdir(newdir, False)

    if gooddir:
        destdir = newdir
        print "setting new destdir"
#------------------------------------------------------------------------------#
def chk_destdir(dir='', exit_on_error=True):

    if dir == '':
        dir = destdir

    # see if we need to create destdir
    if not os.path.isdir(dir):
        prompt = '%s is not a directory. Shall I create it? [Y/n] ' % (dir)
        inp = raw_input(prompt)

        if inp.lower() in ('n', 'no'):
            if exit_on_error:
                sys.exit("quitting")
            else:
                return False
        else:
            try:
                os.makedirs(dir, 0755)
            except:
                pass
            if not os.path.isdir(dir):
                msg = "couldn't create %s" % (dir)
                if exit_on_error:
                    sys.exit(msg)
                else:
                    print msg
                    return False

    # see if we can write to destdir
    if not os.access(dir, os.W_OK):
        msg = "don't have write permissions on %s " % (dir)
        if exit_on_error:
            sys.exit(msg)
        else:
            print msg
            return False

    return True
#------------------------------------------------------------------------------#
def show_verbose_status():
    if verbose:
        print "verbose is currently on"
    else:
        print "verbose is currently off"
#------------------------------------------------------------------------------#
def set_verbose(cmd):
    global verbose

    m = vre.match(cmd)
    on_off = m.group(1)

    if on_off == 'on':
        verbose = True
        print "verbose is now on"
    elif on_off == 'off':
        verbose = False
        print "verbose is now off"
#------------------------------------------------------------------------------#
def show_no_subdirs_status():
    if keep_subdirs:
        print "no-subdirs is currently off"
    else:
        print "no-subdirs is currently on"
#------------------------------------------------------------------------------#
def set_no_subdirs(cmd):
    global keep_subdirs

    m = sdre.match(cmd)
    on_off = m.group(1)

    if on_off == 'on':
        keep_subdirs = False
        print "no-subdirs is now on"
    elif on_off == 'off':
        keep_subdirs == True
        print "no-subdirs is now off"
#------------------------------------------------------------------------------#
def show_usage():
    usg = '''
usage: selcopy -s srcdir -d destdir [-v|--verbose] [--queuefile file]
        [--no-subdirs] [--one-by-one]

    -v, --verbose: increase amount of descriptive output
    --srcdir: alternative for -s
    --destdir alternative for -d
    --queuefile file: copies files listed in file without interactive mode
    --no-subdirs: files all go in the destdir without the original dir structure
    --one-by-one: will cause you to be asked yes or no for each file


    Read the comment at the top of the source for more information.
'''
    sys.exit(usg)
#------------------------------------------------------------------------------#
def show_cmd_help():
    print
    print 'Possible commands:'
    print '-' * 80
    msg = '''    d
        go into the next directory in the list
    dn
        go into subdirectory n
    u
        go up a directory
    pwd
        print working directory
    l
        list current directory contents again; file numbers enclosed in []
        indicate files already in queue
    lq
        list queued files
    qs
        display total filesize of queued files in human readable form
    qsb
        display total filesize of queued files in bytes
    n1 n2 n3 (...nN)
        add files labeled n1, n2, etc to the queue
    a
        add all files in this directory to the queue
    r queuefile
        add all files listed in queuefile to the queue
    w queuefile
        write the queue out to queuefile
    uq f1 f2 f3 (...fN)
        un-queue (remove from queue) files labeled f1, f2, f3, etc. if f starts
        with 'q', files are removed as numbered in the queue itself; otherwise
        files are removed as numbered in the current directory
    copy
        copy the queued files to destdir
    xq
        clear the queue (no copying is done)
    srcdir
        print out the srcdir
    destdir
        print out the destdir
    set destdir dirpath
        change destination directory to dirpath
    verbose
        print out current status of verbosity
    set verbose [on|off]
        turn on or off verbosity
    no-subdirs
        print out current status of no-subdirs
    set no-subdirs [on|off]
        turn on or off no-subdirs
    quit, q
        quit immediately (ignoring any queued files not already copied)

    one-by-one mode commands:
        y
            add current file to queue
        n
            skip current file'''

    
    print msg
    print '-' * 80
#------------------------------------------------------------------------------#
init()
main()