#! /usr/bin/env python

# -*- coding: utf-8 -*-

# vim: set expandtab tabstop=4 shiftwidth=4 :

"""

Convert pictures and videos from a directory to another directory

where the pictures are recompressed and resized and the videos are

converted to FLV, to be displayed by my photos indexer (based on

PrettyPhoto).

Copyright (C) 2008 Aurelien Bompard <aurelien@bompard.org>

This program is free software: you can redistribute it and/or modify

it under the terms of the GNU Affero General Public License as

published by the Free Software Foundation, either version 3 of the

License, or (at your option) any later version.

This program 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 Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License

along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""

import sys, os, optparse, shutil, commands, tempfile, re, stat, locale, datetime

thumb_dir = "_thumbs"

thumb_size = 90

img_size = {"width": 800, "height": 600}

video_thumbs = { "top": "./images/video-top.gif", 

                 "bottom": "./images/video-bottom.gif",

                 "height": 32,

                 "width": 1290 }

use_pil = False # PIL can't preserve EXIF data yet (http://mail.python.org/pipermail/image-sig/2008-January/004787.html)

try:

    import PIL.Image

except ImportError:

    use_pil = False

def processDir(dir, options):

    """ Go through all directories specified on the command line, and convert them to the dest dir """

    dir_clean = getCleanName(dir)

    final_dir = os.path.join(options.dest, dir_clean)

    # Copy the original dir to the temp dir (if not yet done)

    if os.path.exists(final_dir):

        try:

            print "The final directory already exists:", final_dir

            print "Should I go on anyway (no file will be overwritten) ? [Y/n]"

            rep = raw_input()

        except KeyboardInterrupt:

            print "Aborted."

            sys.exit(1)

        if rep.strip() == "n":

            print "Skipped %s !" % dir

            return

    else:

        os.makedirs(final_dir)

        for f in os.listdir(dir):

            if os.path.isdir(os.path.join(dir,f)):

                continue

            shutil.copy(os.path.join(dir,f), 

                        os.path.join(final_dir,f))

            os.chmod(os.path.join(final_dir,f), 

                     stat.S_IRUSR|stat.S_IWUSR|stat.S_IRGRP|stat.S_IROTH) # chmod 644

    # Rename files if needed

    if getattr(options, "prefix", None) is not None:

        renameFilesInDir(final_dir, options)

    # Prepare metadata

    if not os.path.exists(os.path.join(final_dir, "_infos.ini")):

        infofile = open(os.path.join(final_dir, "_infos.ini"), "w")

        dir_name = os.path.basename(dir.rstrip("/"))

        # I usually name my photo dirs this way : "2008-06-16 Joe's birthday"

        # Pre-fill the metadata if that's the format we find

        match = re.match('(20\d\d-\d\d-\d\d) (.*)', dir_name)

        if match:

            title = match.group(2)

            date = match.group(1)

            date = datetime.date(int(date[0:4]),int(date[5:7]),int(date[8:10]))

            date = date.strftime("%A %d %B %Y").capitalize()

        else:

            title = dir_name

            date = ""

        infofile.write("[general]\ntitle = %s\ndate = %s\n" % (title,date))

        infofile.close()

    files = os.listdir(final_dir)

    files.sort()

    print "Recompressing images and building thumbnails from %s..." % dir

    inc=0

    for file in files:

        inc += 1

        if file.lower().startswith("_") or os.path.isdir(os.path.join(final_dir,file)):

            continue # That's either metadata or thumbnails dir

        processFile(file, final_dir, options)

        if not options.verbose: # Nice progress counter

            sys.stdout.write("\r")

            sys.stdout.write("[%s/%s]"%(inc,len(files)))

            sys.stdout.flush()

    if not options.verbose:

        print

def processFile(file, final_dir, options):

    """ Handle each picture or movie """

    if file.lower().endswith(".flv"):

        return # already converted

    if file.lower().endswith(".avi") or file.lower().endswith(".wmv") or file.lower().endswith(".mov"):

        return processVideo(file, final_dir, options)

    if not ( file.lower().endswith(".jpg") or file.lower().endswith(".jpeg") 

             or file.lower().endswith(".gif") or file.lower().endswith(".png") ):

        return # Unknown file format

    filepath = os.path.join(final_dir, file)

    # Portrait or landscape format ?

    portrait = False

    width, height = getSize(filepath)

    if height > width:

        portrait = True

    # Recompress

    recompress(filepath, width, height)

    if options.verbose:

        print "Recompressed %s" % file

    # Prepare thumbnails dir

    if not os.path.exists(os.path.join(final_dir, thumb_dir)):

        os.mkdir(os.path.join(final_dir, thumb_dir))

    # Build thumbnail

    thumb_filename = os.path.join(final_dir, thumb_dir, file)

    if os.path.exists(thumb_filename):

        return

    if use_pil:

        im = PIL.Image.open(filepath)

        if portrait: # resize by height

            new_height = thumb_size

            new_width = int(im.size[0] * thumb_size / float(im.size[1])) # keep proportion

        else: # resize by width

            new_width = thumb_size

            new_height = int(im.size[1] * thumb_size / float(im.size[0])) # keep proportion

        im.thumbnail((new_width,new_height), PIL.Image.ANTIALIAS)

        im.save(thumb_filename)

    else:

        size = thumb_size

        if portrait:

            size = "x%s" % thumb_size

        command = """convert -geometry %s "%s" "%s" """ % (size, filepath, thumb_filename)

        os.system(command)

    if options.verbose:

        print "Built thumbnails for %s..." % file

def processVideo(file, final_dir, options):

    """ Handle movies """

    filepath = os.path.join(final_dir, file)

    thumb_filename = os.path.join(final_dir, thumb_dir, file[:-4]+".jpg")

    # Find size

    command = """ffmpeg -i "%s" 2>&1 | awk '/Video: / {print $6}'""" % filepath

    output = commands.getoutput(command).rstrip(",")

    if output:

        width = int(output.split("x")[0])

        height = int(output.split("x")[1])

    else:

        width = 640

        height = 480

    # Prepare thumbnails dir

    if not os.path.exists(os.path.join(final_dir, thumb_dir)):

        os.mkdir(os.path.join(final_dir, thumb_dir))

    # Make thumbnail

    if not os.path.exists(thumb_filename):

        tmpfd, tmpfile = tempfile.mkstemp(".jpg")

        os.close(tmpfd)

        command= """ffmpeg -vframes 1 -i "%s" -ss 00:00:01.000 -y -f mjpeg "%s" &>/dev/null""" % (filepath, tmpfile) 

        os.system(command)

        if use_pil:

            im = PIL.Image.open(tmpfile)

            im_top = PIL.Image.open(video_thumbs["top"])

            im_bottom = PIL.Image.open(video_thumbs["bottom"])

            im.paste(im_top, (0,0))

            im.paste(im_bottom, (0,height - video_thumbs["height"]) )

            im.thumbnail(thumb_size, PIL.Image.ANTIALIAS)

            im.save(thumb_filename)

        else:

            command = """convert -draw 'image Over 0,0 0,0 "%s"' -draw 'image Over 0,%s 0,0 "%s"' -geometry %s "%s" "%s" """ % \

                         (video_thumbs["top"], height - video_thumbs["height"], video_thumbs["bottom"], thumb_size,

                          tmpfile, thumb_filename)

            os.system(command)

        os.remove(tmpfile)

        if options.verbose:

            print "Built thumbnails for %s..." % file

    # Convert to FLV (open format, readable by gnash)

    if not os.path.exists(filepath[:-4]+".flv"):

        tmpfd, tmpfile = tempfile.mkstemp(".flv")

        os.close(tmpfd)

        command= """ffmpeg -y -i "%s" "%s" &>/dev/null""" % (filepath, tmpfile) 

        os.system(command)

        os.chmod(tmpfile, stat.S_IRUSR|stat.S_IWUSR|stat.S_IRGRP|stat.S_IROTH) # chmod 644

        shutil.move(tmpfile, filepath[:-4]+".flv")

        if options.verbose:

            print "Converted %s to FLV..." % file

    os.remove(filepath)

    # Specify the size in the metadata (needed by shadowbox, so obsolete now)

    #infofile = open(os.path.join(final_dir, "_infos.ini"), "a")

    #infofile.write("[%s]\nwidth = %s\nheight = %s\n"%(file[:-4]+".flv",width,height))

    #infofile.close()

def getSize(file):

    """ Find the size of a picture using PIL or ImageMagick """

    if use_pil:

        im = PIL.Image.open(file)

        return im.size

    else:

        output = commands.getoutput('identify "%s"' % file)

        if file.lower().endswith(".jpg") or file.lower().endswith(".jpeg"):

            properties = output.split(" JPEG ")[1]

        elif file.lower().endswith(".gif"):

            properties = output.split(" GIF ")[1]

        elif file.lower().endswith(".png"):

            properties = output.split(" PNG ")[1]

        else:

            print "Can't identify %s !"%file

            sys.exit(1)

        size = properties.split(" ")[0].split("x")

        return int(size[0]), int(size[1])

def recompress(file, width, height):

    """ Resize pictures and set quality using PIL or ImageMagick """

    if use_pil:

        im = PIL.Image.open(file)

        outim = im

        new_width = width

        new_height = height

        if height > img_size["height"]:

            new_height = img_size["height"]

            new_width = int(width * img_size["height"] / float(height)) # keep proportion

        elif width > img_size["width"]:

            new_width = img_size["width"]

            new_height = int(height * img_size["width"] / float(width)) # keep proportion

        outim = im.resize((new_width, new_height), PIL.Image.ANTIALIAS)

        outim.save(file, quality=75)

    else:

        if height > img_size["height"]:

            os.system("""convert -resize x%s "%s" "/tmp/%s" """ % (img_size["height"], file, os.path.basename(file)) )

            shutil.move("/tmp/%s"%os.path.basename(file), file)

        elif width > img_size["width"]:

            os.system("""convert -resize %s "%s" "/tmp/%s" """ % (img_size["width"], file, os.path.basename(file)) )

            shutil.move("/tmp/%s"%os.path.basename(file), file)

        os.system("""convert -quality 75 "%s" "/tmp/%s" """ % (file, os.path.basename(file)))

        shutil.move("/tmp/%s"%os.path.basename(file), file)

def getCleanName(dirname):

    """ Replace non-ascii chars by their equivalent """

    dir_clean = dirname.replace(" ", "_").replace("'", "").replace('"', '')

    dir_clean = dir_clean.lower()

    # dirty, but efficient ! :)

    mapping = { "é":"e", "è":"e", "ê":"e", "ç":"c", "à":"a", "â":"a", "ô":"o", "î":"i", "ï":"i", }

    for key in mapping:

        dir_clean = dir_clean.replace(key, mapping[key])

    dir_clean = dir_clean.decode("utf-8").encode("ascii", "ignore")

    return dir_clean

def renameFilesInDir(dir, options) :

    """Expects the name of the directory containing the files and

       the prefix to add to each new filename"""

    liste = os.listdir(dir)

    liste.sort()

    nombre = len(liste)

    inc = 1

    print "Renaming files..."

    for fichier in liste:

        if os.path.isdir(os.path.join(dir,fichier)):

            if options.verbose:

                print "Skipping dir %s" % fichier

            continue

        # Extract extention

        dot = fichier.rfind(".")

        if dot == -1 :

            print "ERROR: I can't extract the extention for file %s" % fichier

            sys.exit(1)

        ext = fichier[dot:].lower()

        if ext not in [".jpg", ".jpeg", ".mov", ".avi", ".flv", ".wmv", ".mov"]:

            if options.verbose:

                print "Ignoring %s" % fichier

            continue

        # fill the index with enough zeros on the left

        incLenght = len(str(nombre))

        incStr = str(inc).zfill(incLenght)

        newfile = os.path.join(dir, options.prefix+"_"+incStr+ext)

        oldfile = os.path.join(dir, fichier)

        # while the destination file exists, find another number

        while os.path.exists(newfile):

            if options.verbose:

                print "The file", newfile, "already exists !"

            inc += 1

            incStr = str(inc).zfill(incLenght)

            newfile = os.path.join(dir, options.prefix+'_'+incStr+ext)

            if options.verbose:

                print "Trying", newfile+"..."

        if options.verbose:

            print "Renaming", os.path.basename(oldfile), " --> ", os.path.basename(newfile)

        else:

            sys.stdout.write(".")

            sys.stdout.flush()

        os.rename(oldfile, newfile)

        inc += 1

    if not options.verbose:

        print

def parseOpts():

    usage = "usage: %prog [options] original-folder [original-folder-2 ...]"

    parser = optparse.OptionParser(usage)

    parser.add_option("-p", "--prefix", dest="prefix",

                      help="prefix filenames with this PREFIX", metavar="PREFIX")

    parser.add_option("-d", "--destination", dest="dest",

                      help="the resulting files will be placed in this folder")

    parser.add_option("-v", "--verbose", action="store_true", dest="verbose", default=False)

    (options, args) = parser.parse_args()

    if len(args) == 0:

        parser.error("must specify a directory to process")

    if not options.dest:

        options.dest = "/tmp"

    for dir in args:

        if not os.path.exists(dir) or not os.path.isdir(dir):

            parser.error('The directory %s does not exist.' % dir)

    print

    print "Folder(s) to process:", ", ".join(args)

    print "Destination folder:", options.dest

    # Find subdirs

    dirs = args[:] # copy list object

    total_files = 0

    for dir in dirs:

        for base, subdirs, subfiles in os.walk(dir):

            if subdirs:

                dirs.extend( [os.path.join(base, d) for d in subdirs] )

            total_files += len(subfiles)

    print "Number of files to be processed:", total_files

    if getattr(options, "prefix", None) is not None:

        print "Rename files with prefix:", options.prefix

    print "Is it OK ? (Ctrl-C to abort)"

    raw_input()

    return options, dirs, args

#### MAIN ####

def main(argv) :

    try:

        options, dirs, args = parseOpts()

    except KeyboardInterrupt:

        print "Aborted."

        sys.exit(0)

    # set the correct locale for the names of the days and months

    locale.setlocale(locale.LC_ALL, '')

    dirs.sort()

    for dir in dirs:

        processDir(dir, options)

    # Adapt below to your upload needs

    print "To upload, you may run:"

    upload_cmd = [ "scp -r" ]

    for dir in args:

        upload_cmd.append(os.path.join(options.dest, getCleanName(dir)))

    upload_cmd.append("your.server:your/destination/dir/")

    print " ".join(upload_cmd)

if __name__ == '__main__' : main(sys.argv)