#!/usr/bin/env python

import os
import os.path
import sys
import optparse
import pickle
import subprocess
import random
import signal

# This is where reFLAC stores its current state.  The lack of a
# path means that it will default to storing it in the current
# working directory.  This can be changed by passing --statefile
# to the invocation.

DEFAULT_STATEFILE="reFLAC.state"

# This is how frequent reFLAC will flush its state data; every
# DEFAULT_FREQUENCY songs, it will rewrite the state file.  This
# can be changed by passing --frequency to the invocation.
DEFAULT_FREQUENCY=20

# This is the location of the default FLAC binary.  Because of
# the newness of FLAC 1.1.3, this may not be /usr/bin/flac on
# a typical Linux system.  It can be changed by passing --flac-bin
# to the invocation.
DEFAULT_FLAC_BINARY="/usr/bin/flac"

# This is the percentage of compression below which a WARNING is
# printed.  reFLAC does not recompress files that do not pass a
# flac --test, but there may be some heretofore unknown bug in the
# FLAC->FLAC recompression algorithm.  Since the average shrinkage
# in filesize is ~1%, a loss of 10% is highly suspicious, hence
# the default value of 0.9.  There is no paramter that adjusts this
# particular value, as its need is fairly esoteric.

THRESHOLD_PERCENT=0.9

# Here are the parameters to FLAC that are passed by default
# for re-encoding.  You probably don't want things like
# --replay-gain in here, as it goes file-by-file and will mess
# up your album gain.  But if you want particularly anal-retentive
# encoding options, need to change the way that tukey() and its
# friends are passed due to localisation issues, and so on, you
# can change this list.
REFLAC_PARAMETERS = ['-V', '--best']

# These are global variables, mainly useful for cleanup and storing
# the total amount of data processed by reFLAC across invocations.
# Don't mind the sloppy Python style.
current_reFLAC_file = False
original_sizes = 0
new_sizes = 0

def readState(filename):

   global original_sizes
   global new_sizes

   state = {}
   try:
      fileobj = file(filename, "rb")
      state = pickle.load(fileobj)
      original_sizes = pickle.load(fileobj)
      new_sizes = pickle.load(fileobj)
      fileobj.close()
   except:
      print "Found no preexisting state file."

   return state

def writeState(state, filename):

   global original_sizes
   global new_sizes

   print "reFLAC: Shrank %d bytes to %d bytes." % (original_sizes, new_sizes)
   try:
      fileobj = file(filename, "wb")
      pickle.dump(state, fileobj)
      pickle.dump(original_sizes, fileobj)
      pickle.dump(new_sizes, fileobj)
      return True
   except:
      print "ERROR: Could not write state to file %s!" % filename
      return False

def buildOptionParser():

   usage = "usage: %prog [options] directory"
   parser = optparse.OptionParser(usage=usage)

   parser.add_option("-s", "--statefile", dest="statefile",
    help="Write state to STATEFILE", metavar="STATEFILE",
    default=DEFAULT_STATEFILE, action="store", type="string")
   parser.add_option("-r", "--frequency", dest="frequency",
    help="Flush state every FREQ files", metavar="FREQ",
    action="store", type="int", default=DEFAULT_FREQUENCY)
   parser.add_option("-f", "--force", action="store_true",
    help="Continue even if the state file is unwritable",
    default=False, dest="force")
   parser.add_option("-F", "--flac-bin", action="store",
    type="string", help="Use BIN as executable", metavar="BIN",
    default=DEFAULT_FLAC_BINARY, dest="flac")
   parser.add_option("-p", "--pretend", action="store_true",
    dest="pretend", help='Don\'t actually make new FLACs',
    default=False)
   parser.add_option("-v", "--verbose", action="store_true",
    dest="verbose", help='Be chatty',
    default=False)

   return parser

def reFLAC(path, state, options):

   global current_reFLAC_file
   global original_sizes
   global new_sizes

   file_count = 0
   for root, dirs, files in os.walk(path):
      for name in files:

         # We only care about files that end in ".flac," because we're too
         # lazy to do anything fancy with mimetypes.  We also want canonical
         # paths, because this program could be invoked all sorts of ways.
         if len(name) > 5 and name[-5:] == ".flac":
            full_filename = os.path.realpath(os.path.join(root, name))
            if options.verbose:
               print "reFLAC: [%s] Considering ..." % full_filename
               sys.stdout.flush()

            # Record any file we haven't seen before in the global state.
            if full_filename not in state:
               state[full_filename] = False

            # If state tells us we've already processed this file, no need
            # to waste CPU doing it again.
            if state[full_filename]:
               if options.verbose:
                  print "reFLAC: [%s] Already handled; skipping." % full_filename
                  sys.stdout.flush()

            # "QQQreFLACQQQ.flac" is our temporary filename, and is highly,
            # highly unlikely to be the name of any file generated by anyone
            # else.  Skip files named this and warn the user.
            elif len(full_filename) > 17 and name[-17:] == "QQQreFLACQQQ.flac":
               print "reFLAC: [%s] reFLAC file; skipping.  You should probably delete this." % full_filename
               sys.stdout.flush()
            else:
               if options.verbose:
                  print "reFLAC: [%s] Testing ..." % full_filename
                  sys.stdout.flush()
               file_count += 1

               # Test the FLAC via flac -t.  We're not going to process it if
               # this fails, to work around the 1.1.3 bug that dies on
               # FLAC-to-FLAC conversion if there's an error in the FLAC, not
               # to mention it's probably a bad idea anyway.
               ret = subprocess.call([options.flac, "-s", "-t", full_filename],
                stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                stderr=subprocess.PIPE)
               if ret:
                  print "ERROR: File %s is not a valid FLAC; skipping." % full_filename
                  sys.stdout.flush()
               elif not options.pretend:
                  if options.verbose:
                     print "reFLAC: [%s] Re-encoding ..." % full_filename
                     sys.stdout.flush()
                  new_filename = os.path.realpath(os.path.join(root,
                   "QQQreFLACQQQ.flac"))
                  current_reFLAC_file = new_filename

                  # This is where we actually do the reencoding.  We build
                  # the command-line from the REFLAC_PARAMETERS option
                  # above, the FLAC executable location, and the source
                  # and destination filenames.
                  options_list = [options.flac, "-s"]
                  options_list.extend(REFLAC_PARAMETERS)
                  options_list.extend([full_filename, "-o", new_filename])

                  # If it returns anything other than a 0, something went
                  # wrong, and we're not going to keep the new file.
                  ret = subprocess.call(options_list, stdin=subprocess.PIPE,
                   stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                  if ret:
                     print "ERROR: Could not reencode FLAC %s; skipping." % full_filename
                     try:
                        print "[reFLAC] Cleaning up %s ..." % new_filename
                        os.unlink(new_filename)
                     except:
                        pass
                     sys.stdout.flush()
                  else:
                     original_size = os.path.getsize(full_filename)
                     new_size = os.path.getsize(new_filename)

                     # We're not going to replace files unless they're smaller;
                     # it's stupid to switch out for a larger file.
                     if new_size < original_size:
                        print "reFLAC: [%s] %d -> %d; switching." % (full_filename,
                         original_size, new_size)
                        if new_size < original_size * THRESHOLD_PERCENT:
                           print "WARNING: The previous file shrank to less than %s%% its original size!  This may be an error." % repr (THRESHOLD_PERCENT * 100)
                        sys.stdout.flush()
                        original_sizes += original_size
                        new_sizes += new_size
                        os.unlink(full_filename)
                        os.rename(new_filename, full_filename)
                        current_reFLAC_file = False
                     else:
                        print "reFLAC: [%s] No savings." % (full_filename)
                        sys.stdout.flush()
                        os.unlink(new_filename)
                     state[full_filename] = True
         if file_count >= options.frequency:
            file_count = 0
            if options.verbose:
               print "reFLAC: Writing state."
               sys.stdout.flush()
            writeState(state, options.statefile)

   if options.verbose:
      print "reFLAC: Writing state."
      sys.stdout.flush()
   writeState(state, options.statefile)
   sys.stdout.flush()


def cleanup(state, statefile):

   global current_reFLAC_file

   if current_reFLAC_file:
      print "[reFLAC] Cleaning up %s ..." % current_reFLAC_file
      os.unlink(current_reFLAC_file)
   print "[reFLAC] Saving state ..."
   writeState(state, statefile)

# This allows us to clean up if someone sends a standard 'kill'
# call, which is nice for long-running processes.
def sighandler(signum, frame):

   global options
   global state

   print "[reFLAC] Received signal %d." % signum
   cleanup(state, options.statefile)

   
if "__main__" == __name__:

   signal.signal(signal.SIGTERM, sighandler)

   optparser = buildOptionParser()
   (options, args) = optparser.parse_args()
   if len (args) != 1 or not os.path.isdir(args[0]):
      print "ERROR: Must give a directory as the argument."
      sys.exit(1)
   state = readState(options.statefile)
   if not writeState(state, options.statefile) and not options.force:
      print "ERROR: Could not rewrite statefile."
      sys.exit(1)
   try:
      reFLAC(args[0], state, options)
   except:
      cleanup(state, options.statefile)