#! /usr/bin/python

import os, sys, stat, getopt, re, traceback, random, statvfs, re

__version__ = "0.2"
BUF_SIZE = 10240

def showVersion():
  print 'rimshuffle v'+__version__+', by Rimon Barr:',
  print 'Random file copier.'

def showUsage():
  print
  showVersion()
  print
  print 'Usage: rimshuffle src dst'
  print '       rimshuffle -h | -v'
  print 
  print '  -h, -?, --help      display this help information'
  print '  -v, --version       display version'
  print '  -r, --recurse       recurse for files in src directory'
  print '  -l, --limit         size limit'
  print '  -n, --files         file limit'
  print '  -d, --delete        delete first'
  print
  print 'Send comments, suggestions and bug reports to <rimon@acm.org>.'
  print

def freespace(path):
  if os.name=='nt':
    import win32file
    return win32file.GetDiskFreeSpaceEx(path)[0]
  else:
    s = os.statvfs(path)
    return s[statvfs.F_BAVAIL] * long(s[statvfs.F_BSIZE])

def filesize(path):
  st=os.lstat(path)
  return long(st[stat.ST_SIZE])

def filecopy(srcpath, dstpath):
  fsrc = fdst = None
  erase = 0
  try:
    try:
      fsrc, fdst = open(srcpath, 'rb'), open(dstpath, 'wb')
      erase = 1
      x = fsrc.read(BUF_SIZE)
      while x: fdst.write(x); x=fsrc.read(BUF_SIZE)
      fdst.flush()
      erase = 0
    finally:
      try: 
        if fsrc: fsrc.close()
      except: pass
      try: 
        if fdst: fdst.close()
      except: pass
      if erase: os.remove(dstpath)
  except IOError: return -1
  return 1

def parsesize(s):
  m = re.match(r'(\d+)(.*)', s)
  if not m: return -1
  size = int(m.group(1))
  unit = m.group(2).strip().upper()
  if unit == '': return size
  if unit == 'K': return size * 1024L
  if unit == 'M': return size * 1024L * 1024L
  return -1

def usageError():
  print 'rimshuffle: command syntax error'
  print 'Try `rimshuffle --help\' for more information.'

# format file size (bytes)
def formatBytes(bytes):
  if bytes<=99999: return '%8.2f' % bytes
  bytes = bytes / 1024.0
  if bytes<=9999: return '%8.2fKi' % bytes
  bytes = bytes / 1024.0
  if bytes<=9999: return '%8.2fMi' % bytes
  bytes = bytes / 1024.0
  return '%8.2fGi' % bytes

def getFiles(src, recurse=1):
  filelist = []
  for root, dirs, files in os.walk(src):
    files = [os.path.join(root, f) for f in files]
    if not recurse: return files
    filelist = filelist + files
  return filelist

def rimshuffle():
  # parse options
  try: opts, args = getopt.getopt(sys.argv[1:], 'hv?rl:n:',
      ['help', 'version', 'recurse', 'limit=', 'files='])
  except getopt.GetoptError: usageError(); return
  # read src and dst directories
  if len(args)<2: showUsage(); return
  src, dst = args[0], args[1]
  # option defaults
  recurse, maxfile, space, deletefirst = False, 0, 0, False
  # process options
  for o, a in opts:
    if o in ("-h", "--help", "-?"): showUsage(); return
    if o in ("-v", "--version"): showVersion(); return
    if o in ("-r", "--recurse"): recurse = True
    if o in ("-l", "--limit"): space = max(parsesize(a), 0)
    if o in ("-n", "--files"): maxfile = int(a)
    if o in ('-d', '--delete'): deletefirst = True
  # delete files
  if deletefirst: raise NotImplementedError
  # perform the copy
  if space: space = min(space, freespace(dst))
  else:     space = freespace(dst)
  if not os.path.isdir(src): print 'src must be an existing directory'; return 1
  if not os.path.isdir(dst): print 'dst must be an existing directory'; return 1
  files = getFiles(src, recurse)
  random.shuffle(files)
  filecount = len(getFiles(dst, 0))
  spaceleft, count = space, 0
  while files:
    srcfile = files[0]
    del files[0]
    if not os.path.exists(srcfile): continue
    size = filesize(srcfile)
    if spaceleft < size: break
    if maxfile and filecount >= maxfile: break
    dstfile = os.path.join(dst, os.path.split(srcfile)[1])
    if os.path.exists(dstfile): continue
    if not filecopy(srcfile, dstfile):
      print 'failed to copy: ', srcfile
      break
    spaceleft, filecount = spaceleft-size, filecount+1
    if maxfile: print '%4d' % (maxfile-filecount),
    print '%5s' % formatBytes(spaceleft),
    print '- %s' % os.path.abspath(srcfile)[len(os.path.abspath(src))+1:]
    count += 1
  print count, "files copied."

try:
  if __name__=='__main__': success = rimshuffle()
  sys.exit(success)
except IOError: pass
except KeyboardInterrupt: print 'Break!'