#!/usr/bin/python

# reshuffle -- refill an ipod shuffle

# Copyright 2008 Peter Maydell <pmaydell@chiark.greenend.org.uk>
#
#    reshuffle is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 3 of the License, or
#    (at your option) any later version.
#
#    reshuffle is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# iTunes, iPod and Shuffle are trademarks of Apple.
# This product is not supported/written/published by Apple!

# Summary of operation:

# 1. Find (possibly mounting) the shuffle
# 2. Find tracks on the shuffle with a playcount > 0
# 3. Delete the tracks from the shuffle
# 4. Pick a track at random from the computer from
#    the list of tracks which haven't been marked 'played'
#    If we can't find any such track, then remove all
#    the 'played' marks and try again.
# 5. Put it on the ipod, encoding it as we do so,
#    and including in the track's 'userdata' the
#    on-computer filename
# 6. Mark the track as 'played' on the computer
#    (we have to do this as we put the track on the ipod
#    otherwise a subsequent top-up of the ipod might
#    select it again)
# 7. Repeat 4-6 until shuffle is full
# 8. Cleanly unmount the shuffle

# On disk data:
# Each subdirectory containing music files may also
# have a file 'reshuffle.played' which is a text file
# which lists all the files in that directory which
# have been played in this run. If no file exists then
# this is the same as having an empty file.
# There may also be a file 'reshuffle.ignore'; contents
# are unimportant. Any subdirectory containing this file
# will be ignored when scanning for tracks to add to
# the ipod. (This script doesn't create such files; the
# idea is that the user manually creates them to stop
# us looking in directories containing things not wanted.)

# FIXME it would assist debugging if we set at least some
# of the artist/title info in the db

# FIXME we should probably have multiple verbosity
# levels: debug/verbose/silent

# FIXME it would be nice if we didn't cause the shuffle
# to forget where it was in the track it was currently playing

# FIXME only handling FLAC is fine for us but we could
# be more generalised

import gpod
import optparse
import os
import os.path
import random
import re
import subprocess
import tempfile
import glob

class Global:
    """Generic container for shared variables."""
    pass

IGNORE_FILE="reshuffle.ignore"
PLAYED_FILE="reshuffle.played"

# log levels:
# debugging, if verbose, if verbose or if dry-run, print even if quiet
LOG_DEBUG=3
LOG_VERBOSE=2
LOG_DRYRUN=1
LOG_QUIET=0

def log(loglevel,string):
    if Global.loglevel >= loglevel:
        print string

def mark_played_on_disk(filename):
    log(LOG_DRYRUN, "Marking %s as played on disk" % filename)
    if Global.dry_run:
        return
    dir = os.path.dirname(filename)
    f = open(os.path.join(dir, PLAYED_FILE), "a")
    f.write(os.path.basename(filename) + "\n")
    f.close()

def remove_dirs_to_ignore(root,dirlist):
    # Update dirlist in-place to remove names of
    # directories which contain an IGNORE_FILE
    # The in-placeness is important as this is how
    # os.walk works.
    # We go backwards so that one deletion doesn't
    # mess up the index number for a later deletion.
    for i in reversed(range(len(dirlist))):
        if os.path.isfile(os.path.join(root, dirlist[i], IGNORE_FILE)):
            log(LOG_DEBUG, "Ignoring directory %s" % dirlist[i])
            del dirlist[i]

def delete_played_file(pfile):
    log(LOG_DRYRUN, "Deleting file: %s" % pfile)
    if Global.dry_run:
        return
    os.remove(pfile)

def clear_all_played_marks(directory):
    # Clear all the 'played' marks from directory
    log(LOG_DRYRUN, "Clearing all the 'played' marks")
    for root, dirs, files in os.walk(directory):
        remove_dirs_to_ignore(root,dirs)
        if PLAYED_FILE in files:
            delete_played_file(os.path.join(root,PLAYED_FILE))

def is_music_file(filename):
    # Return true if filename is a music file
    return filename.endswith(".flac")

def get_unplayed_tracks(directory):
    # Return a shuffled list of all the tracks which
    # haven't been played yet
    log(LOG_VERBOSE, "Searching for unplayed tracks...")
    tracks = []
    for root, dirs, files in os.walk(directory):
        remove_dirs_to_ignore(root,dirs)
        played = []
        if PLAYED_FILE in files:
            f = open(os.path.join(root, PLAYED_FILE))
            try:
                played = [line.rstrip('\n') for line in f.readlines()]
                log(LOG_DEBUG, "Got played tracks for %s: %r" % (root, played))
            finally:
                f.close()

        files = [f for f in files if f not in played and is_music_file(f)]
        tracks.extend([os.path.join(root,f) for f in files])
    log(LOG_DEBUG, "Finished scanning for tracks")
    random.shuffle(tracks)
    return tracks

def convert_file_to_mp3(infile,mp3file):
    # Convert a FLAC file to mp3
    log(LOG_DEBUG, "Converting to MP3...")
    p1 = subprocess.Popen(['flac', '-d', '-s', '-c', infile], stdout=subprocess.PIPE)
    p2 = subprocess.Popen(['lame', '--silent', '-V', '0', '-o', '-', mp3file], stdin=p1.stdout)
    if p2.wait() or p1.wait():
        # oops, something went wrong
        log(LOG_QUIET, "Error in conversion process!")
        return False
    return True

def set_trackdata(track,filename):
    # Set any relevant metadata for this track if we can
    # Just save the original filename into the db
    # (my music doesn't have FLAC tags and splitting the
    # filename into artist/album/track seems rather specific
    # to my personal setup)
    track['userdata']['reshufflefile'] = filename
    
def add_track_to_ipod(db,filename):
    # Add the specified track to the ipod.
    # Return True if we managed this, False if
    # we failed (probably due to filling the ipod)
    log(LOG_DRYRUN, "Adding track %s to ipod" % filename)
    if Global.dry_run:
        return True
    mp3file = tempfile.NamedTemporaryFile(prefix='reshuffle', suffix='.mp3')
    try:
        if not convert_file_to_mp3(filename, mp3file.name):
            return False
        track = db.new_Track(filename=mp3file.name)
        set_trackdata(track, filename)
        try:
            track.copy_to_ipod()
        except gpod.TrackException, e:
            # Important that we remove this failed track otherwise
            # we are left with a duff entry in the db (which causes
            # complaints later)
            log(LOG_VERBOSE, "Copy failed, backing it out")
            db.remove(track)
            return False
        return True
    finally:
        mp3file.close()

def get_trackstr(track):
    trackstr = "%r" % track
    if track['userdata']:
        try:
            trackstr = track['userdata']['reshufflefile']
        except KeyError:
            pass
    return trackstr

def get_mp3_filename(track):
    # NB: this returns the filename with
    # (a) the mountpoint prefix stripped
    # (b) directory separators replaced with ':'
    name = None
    if track['userdata']:
        try:
            name = track['userdata']['filename_ipod']
        except KeyError:
            pass
    return name

def remove_track_from_ipod(db,track):
    log(LOG_DRYRUN, "Removing from ipod: %s" % get_trackstr(track))
    if Global.dry_run:
        return
    db.remove(track)

def remove_old_tracks(db,directory):
    # Remove any old (played) tracks from the ipod
    log(LOG_DRYRUN, "Searching ipod for played tracks...")
    # Note that we have to construct the tracklist first because you
    # can't remove items from the db while you're iterating over it
    tracklist = [track for track in db if track['playcount'] != 0 or track['skipcount'] != 0]
    for track in tracklist:
        remove_track_from_ipod(db,track)

def remove_early_dups(db,directory):
    # Remove tracks from the start of the ipod which are
    # present multiple times. This is a hacky workaround
    # for gpod not spotting skipped tracks. The idea is
    # that there are a few tracks we never listen to (speech
    # ones, mostly), and they pile up at the front of the
    # ipod. So if any track is present multiple times in
    # the first ten tracks we delete it from the ipod.
    log(LOG_DRYRUN, "Searching ipod for early duplicate tracks...")
    seen = {}
    for track in db[0:9]:
        trackstr = get_trackstr(track)
        if not trackstr in seen:
            seen[trackstr] = []
        seen[trackstr].append(track)
    for tracklist in seen.values():
        if len(tracklist) > 1:
            for track in tracklist:
                 remove_track_from_ipod(db,track) 

def garbage_collect(db,mount):
    # Remove any mp3s from the ipod which don't correspond
    # to entries in the tracklist. These shouldn't exist
    # but if we crashed or something they might get created.
    log(LOG_DRYRUN, "Garbage collecting unreferenced mp3s on ipod...")
    known = {}
    for track in db:
        name = get_mp3_filename(track)
        if name:
            known[name] = True
    for f in glob.iglob(os.path.join(mount,"iPod_Control","Music","*","*.mp3")):
        fname = f.replace(mount, '').replace(os.path.sep, ':')
        if not fname in known:
            delete_played_file(f)

def add_new_tracks(db,directory,addcount):
    # Pick new tracks at random and add them to the ipod
    seen_all = False
    available_tracks = get_unplayed_tracks(directory)
    while True:
        if addcount != None:
            if addcount == 0:
                break
            addcount = addcount - 1
        if available_tracks == []:
            if seen_all:
                log(LOG_DRYRUN, "Completely out of tracks to add")
                return
            seen_all = True
            clear_all_played_marks(directory)
            available_tracks = get_unplayed_tracks(directory)
            continue
        track = available_tracks.pop()
        if not add_track_to_ipod(db,track):
            # run out of space on ipod for more tracks
            break
        mark_played_on_disk(track)

def do_db_close(db):
    if not Global.dry_run:
        log(LOG_VERBOSE, "Writing shuffle database file...")
        if not gpod.itdb_shuffle_write(db._itdb,None):
            raise gpod.DatabaseException("Unable to save shuffle db")

    log(LOG_DRYRUN, "Closing database")
    db.close()

def update_and_close_db(db, mount):
    # We don't do a copy_delayed_files() because
    # we guarantee to have done the copies one at a time
    # earlier.

    # We have to write the shuffle's database file by hand
    # using the low level API (this is a bug in the python
    # bindings which is fixed in upstream SVN but has not yet
    # made it into ubuntu.)
    log(LOG_DRYRUN, "Writing shuffle database")
    try:
        do_db_close(db)
    except Exception, e:
        # If we couldn't write the db, try deleting the last track
        # and trying again. If that fails then just barf.
        db.remove(db[-1])
        do_db_close(db)

    # Remove the stats db (playcounts have been merged into the full database
    # and besides track numbers etc will all be out of sync)
    # (do this even if doing a dry run because libgpod will have updated
    # playcounts in iTunesDB even if we did nothing to the db ourselves)
    statsfile = os.path.join(mount, 'iPod_Control', 'iTunes', 'iTunesStats')
    if os.path.exists(statsfile):
        os.remove(statsfile)

def plausible_mounted_ipod_mountpoint(mount):
    return os.path.isdir(os.path.join(mount, "iPod_Control"))

def mountpoint_in_fstab(mount):
    pat = re.compile(r'[^#]*\s+' + mount + r'\s+')
    try:
        f = open('/etc/fstab', 'r')
    except Exception, e:
        return False
    try:
        for l in f:
            if re.search(pat, l):
                return True;
        return False
    finally:
        f.close()

def try_to_mount(mount):
    subprocess.call(['mount', mount])

def plausible_ipod_mountpoint(mount):
    # Return true if mountpoint looks plausible. If it looks
    # like it might be an unmounted filesystem, try mounting it.
    if plausible_mounted_ipod_mountpoint(mount):
        return True
    if not mountpoint_in_fstab(mount):
        return False
    try_to_mount(mount)
    return plausible_mounted_ipod_mountpoint(mount)

def guess_ipod_mountpoint():
    log(LOG_VERBOSE, "iPod mount point not specified, guessing...")
    for guess in ["/mnt/ipod", "/media/IPOD"]:
        log(LOG_VERBOSE, "Trying %s..." % guess)
        if plausible_ipod_mountpoint(guess):
            log(LOG_VERBOSE, "Yes, %s looks plausible" % guess)
            return guess
    return None

def main():
    oparser = optparse.OptionParser(usage = 'usage: %prog [options] musicdir')
    oparser.add_option('-m', '--mount',
                       dest='mount',
                       help='Specify mountpoint of iPod')
    oparser.add_option('-v', '--verbose',
                       dest='verbose',
                       default=False,
                       action="store_true",
                       help='Be verbose')
    oparser.add_option('--debug',
                       dest='debug',
                       default=False,
                       action="store_true",
                       help='Print debug progress messages (implies -v)')
    oparser.add_option('--dry-run',
                       dest='dry_run',
                       default=False,
                       action="store_true",
                       help='Just print what we would do, without writing to ipod or disk')
    oparser.add_option('--no-delete',
                       dest='no_delete',
                       default=False,
                       action="store_true",
                       help="Don't delete any tracks from the iPod, just try to add new ones")
    oparser.add_option('--add-count',
                       dest='addcount', type='int',
                       help='Add only this many tracks (default is as many as possible)')
    oparser.add_option('--no-unmount',
                       dest='unmount',
                       default=True,
                       action='store_false',
                       help="Don't unmount the iPod when done")
    oparser.add_option('--garbage-collect',
                       dest='garbage_collect',
                       default=False,
                       action='store_true',
                       help="Delete mp3s which aren't in the playlist")
                       
    opts,args = oparser.parse_args()
    if len(args) > 1:
        oparser.error("Too many arguments")
    elif len(args) == 0:
        oparser.error("Music directory not specified")
    directory = args[0]

    Global.dry_run = opts.dry_run

    Global.loglevel = LOG_QUIET
    if Global.dry_run:
        Global.loglevel = LOG_DRYRUN
    if opts.verbose:
        Global.loglevel = LOG_VERBOSE
    if opts.debug:
        Global.loglevel = LOG_DEBUG

    if not opts.mount:
        opts.mount = guess_ipod_mountpoint()
        if not opts.mount:
            oparser.error("Mount point not specified and not guessable")
    if not plausible_ipod_mountpoint(opts.mount):
        oparser.error("Mount point '%s' doesn't look like it has an ipod mounted" % opts.mount)

    log(LOG_DRYRUN, "Opening database for iPod mounted at %s" % opts.mount)
    db = gpod.Database(opts.mount)
    if opts.garbage_collect:
        garbage_collect(db,opts.mount)
    if not opts.no_delete:
        remove_early_dups(db,directory)
        remove_old_tracks(db,directory)
    add_new_tracks(db,directory,opts.addcount)

    update_and_close_db(db,opts.mount)

    if opts.unmount:
        log(LOG_DRYRUN, "Unmounting iPod...")
        if not Global.dry_run:
            if subprocess.call(['umount', opts.mount]):
                print "warning: unmount failed"
            subprocess.call(['sync'])
    
    log(LOG_DRYRUN, "Done.")

if __name__ == '__main__':
    main()