#!/usr/bin/env python

"""Simple utility to resize multiple pictures recursively and concurrently.

In addition to resizing the pictures, this utility will also write a JSON metadata file; for every resized file, the src and dst paths as well as the src and dst md5 will be recorded.
"""

from __future__ import print_function

from Queue import Queue
from subprocess import Popen, PIPE
from threading import Thread
import argparse
import errno
import json
import multiprocessing
import os
import re
import sys

devnull = open('/dev/null', 'w')

class Message(object):
    pass

class FileMessage(Message):
    def __init__(self, src, dst):
        self.src = src
        self.dst = dst

class Success(FileMessage):
    def __init__(self, src, dst, src_md5, dst_md5):
        super(Success, self).__init__(src, dst)
        self.src_md5 = src_md5
        self.dst_md5 = dst_md5
    def __str__(self):
        return '%s (%s) -> %s (%s)' % (self.src, self.src_md5[:6], self.dst, self.dst_md5[:6])
    
class Failure(FileMessage):
    def __init__(self, src, dst, error):
        super(Failure, self).__init__(src, dst)
        self.error = error
    def __str__(self):
        return '%s -> %s: %s' % (self.src, self.dst, self.error)

class Estimate(Message):
    def __init__(self, total):
        self.total = total

class ToolError(Exception):
    def __init__(self, tool_stderr):
        self.tool_stderr = tool_stderr
    def __str__(self):
        return self.tool_stderr

def parse_arguments(argv):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--parallel', default=None, type=int)
    parser.add_argument('--size', default='1920x1200')
    parser.add_argument('--overwrite', default=False, action='store_true')
    parser.add_argument('--allow-existing', default=False, action='store_true')
    parser.add_argument('--allow-whitelist', default='\.(jpg)$')
    parser.add_argument('--allow-whitelist-case-sensitive', default=False, action='store_true')
    parser.add_argument('--metadata-path', default='.metadata.json')
    parser.add_argument('--src-md5-links', default=None)
    parser.add_argument('--src-md5-hardlinks', default=False, action='store_true')
    parser.add_argument('src')
    parser.add_argument('dst')
    options = parser.parse_args(argv[1:])
    if options.parallel is None:
        try:
            options.parallel = max(1, int(multiprocessing.cpu_count()/2))
        except NotImplementedError:
            options.parallel = 1
    if options.allow_whitelist_case_sensitive:
        options.allow_whitelist = re.compile(options.allow_whitelist)
    else:
        options.allow_whitelist = re.compile(options.allow_whitelist, re.IGNORECASE)
    if options.metadata_path and not os.path.isabs(options.metadata_path):
        options.metadata_path = os.path.join(options.dst, options.metadata_path)
    return options

def run_tool(args):
    process = Popen(args, stdin=devnull, stdout=PIPE, stderr=PIPE)
    out, err = process.communicate()
    if process.wait() != 0:
        raise ToolError(err.strip())
    return out

def convert_image(options, src, dst):
    run_tool(['gm', 'convert', '-size', options.size, src, '-resize', options.size, dst])

def get_file_md5(path):
    return run_tool(['md5sum', path]).split()[0]

def start_thread(func, *args, **kwargs):
    t = Thread(target=func, args=args, kwargs=kwargs)
    t.daemon = True
    t.start()
    return t

def worker(options, work_queue, report_queue):
    while True:
        src, dst = work_queue.get()
        try:
            convert_image(options, src, dst)
            src_md5 = get_file_md5(src) if options.metadata_path else None
            if options.src_md5_links:
                action = os.link if options.src_md5_hardlinks else os.symlink
                try:
                    action(src, os.path.join(options.src_md5_links, src_md5))
                except OSError, error:
                    if errno.EEXIST != error.errno:
                        raise
            dst_md5 = get_file_md5(dst) if options.metadata_path else None
        except Exception, error:
            report_queue.put(Failure(src, dst, error))
        else:
            report_queue.put(Success(src, dst, src_md5, dst_md5))
        work_queue.task_done()

def reporter(options, report_queue, metadata):
    success = 0
    failure = 0
    pv = None
    while True:
        msg = report_queue.get()
        if isinstance(msg, FileMessage) and pv is not None:
            pv.stdin.write(' ')
        if isinstance(msg, Success):
            success += 1
            metadata.append({'src': msg.src, 'dst': msg.dst, 'src_md5': msg.src_md5, 'dst_md5': msg.dst_md5})
        if isinstance(msg, Failure):
            failure += 1
            print(msg)
        if isinstance(msg, Estimate):
            assert pv is None, 'unexpected second estimate'
            pv = Popen(['pv', '--progress', '--timer', '--eta', '--average-rate', '--bytes',
                        '--size', str(msg.total - success - failure)], stdin=PIPE, stdout=devnull)
        report_queue.task_done()

def main(options):
    if not options.allow_existing and os.path.isdir(options.dst):
        raise SystemExit('refusing to use existing directory %s' % (options.dst,))

    work_queue = Queue()
    report_queue = Queue()

    if options.src_md5_links and not os.path.isdir(options.src_md5_links):
        os.makedirs(options.src_md5_links)

    metadata = []
    for index in range(options.parallel):
        start_thread(worker, options, work_queue, report_queue)
    start_thread(reporter, options, report_queue, metadata)

    total = 0
    for directory, subdirectories, filenames in os.walk(options.src):
        src_dir = os.path.join(options.src, directory)
        dst_dir = os.path.join(options.dst, os.path.relpath(src_dir, options.src))
        for filename in filenames:
            if not options.allow_whitelist.search(filename):
                continue
            if not os.path.isdir(dst_dir):
                os.makedirs(dst_dir)
            src = os.path.join(src_dir, filename)
            dst = os.path.join(dst_dir, filename)
            if os.path.exists(dst) and not options.overwrite:
                continue
            total += 1
            work_queue.put((src, dst))
    report_queue.put(Estimate(total))

    try:
        work_queue.join()
        report_queue.join()
    except KeyboardInterrupt:
        pass

    if options.metadata_path:
        with open(options.metadata_path, 'w') as handle:
            json.dump(metadata, handle, indent=4)

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