#!/usr/bin/env python2
"reimage: downscale and/or re-encode large JPEG images"
__help__ = """
Options:
  -h, --help               show this help screen
  -r RES, --res RES        set maximum resolution to RES
  -s SIZE, --size SIZE     enable size limit mode and set size limit to SIZE
  -q QUAL, --quality QUAL  enable constant-quality mode, disable size limit
                           and set JPEG encoding quality to QUAL (1..100)
  -o FILE, --output FILE   write output to FILE instead of replacing the input
  -f, --force              force re-encoding even if no change is needed
  -p, --pretend            don't write any files to disk
  -v, --verbose            enable verbose operation

The resolution can be either a width-by-height expression like `1024x768' or a
single number, in which case the width and height limit will be set to the same
value.
The size limit can be specified in bytes, KiB if a `k' is appended or MiB if an
`M' is appended. When a size limit is specified, all output files will be
encoded such that they are just below this size limit. All input files above
this limit will be re-encoded.
"""
import sys, os, stat, Image, ImageFile, cStringIO
ImageFile.MAXBLOCK = 4*1024*1024

MaxRes = None
MaxSize = None
Quality = 75
ForceRepack = False
OutputFile = None
HadOutputFile = False
EnableWrite = True
Verbosity = 0
def v(level): return Verbosity >= level


################################################################################


def bytes(n):
    s = str(n)
    p = len(s) - 3
    while p > 0:
        s = s[:p] + ',' + s[p:]
        p -= 3
    return s


def encode(img, quality):
    f = cStringIO.StringIO()
    img.save(f, "JPEG", quality=quality, optimize=1)
    return f.getvalue()


def Process(infile, outfile):
    if v(2): print "processing", infile, "->", outfile
    try:
        insize = os.stat(infile)[stat.ST_SIZE]
    except OSError, e:
        print >>sys.stderr, "Error reading input file `%s':" % infile, e
        return
    try:
        img = Image.open(infile)
        oldres = img.size
    except IOError, e:
        print >>sys.stderr, "Error reading input file `%s':" % infile, e
        return

    need_repack = ForceRepack or (MaxSize and (insize > MaxSize))
    if MaxRes:
        sx, sy = oldres
        if (sx > MaxRes[0]) or (sy > MaxRes[1]):
            sx = MaxRes[0]
            sy = (oldres[1] * sx + oldres[0] / 2) / oldres[0]
            if sy > MaxRes[1]:
                sy = MaxRes[1]
                sx = (oldres[0] * sy + oldres[1] / 2) / oldres[1]
            if v(2):
                print "resizing %dx%d -> %dx%d" % (oldres[0], oldres[1], sx, sy)
            try:
                img.thumbnail((sx, sy), Image.ANTIALIAS)
            except IOError, e:
                print >>sys.stderr, "Error reading input file `%s':" % infile, e
                return
            need_repack = True
    newres = img.size

    if not need_repack:
        del img
        if infile != outfile:
            try:
                f = open(infile, "rb")
                data = f.read()
                f.close()
            except IOError, e:
                print >>sys.stderr, "Error reading input file `%s':" % infile, e
                return
            try:
                if EnableWrite:
                    f = open(outfile, "wb")
                    f.write(data)
                    f.close()
            except IOError, e:
                print >>sys.stderr, "Error writing output file `%s':" % outfile, e
                return
            if v(1):
                print "%s (%dx%d, %s bytes) -> %s (unchanged)" % (infile, oldres[0], oldres[1], bytes(insize), outfile)
        elif v(1):
            print "%s (%dx%d, %s bytes) -- no change" % (infile, oldres[0], oldres[1], bytes(insize))
        return

    if not(outfile.rsplit('.', 1)[-1].lower() in ('jpg', 'jpeg', 'jpe')):
        outsize = '???'
        try:
            if EnableWrite:
                img.save(outfile)
        except IOError, e:
            print >>sys.stderr, "Error writing output file `%s':" % outfile, e
            return
        try:
            outsize = os.stat(outfile)[stat.ST_SIZE]
        except OSError, e:
            print >>sys.stderr, "Error writing output file `%s':" % outfile, e
        if v(1):
            print "%s (%dx%d, %s bytes) -> %s (%dx%d, %s bytes)" % (infile, oldres[0], oldres[1], bytes(insize), outfile, newres[0], newres[1], bytes(outsize))

    if Quality:
        data = encode(img, Quality)
    else:
        a = 1
        b = 95
        data = ""
        while a < b:
            c = (a + b) / 2
            check = encode(img, c)
            if v(2):
                print "quality %d -> %s bytes" % (c, bytes(len(check)))
            size = len(check)
            if size <= MaxSize:
                if size > len(data):
                    data = check
                a = c + 1
            else:
                b = c
    del img

    try:
        if EnableWrite:
            f = open(outfile, "wb")
            f.write(data)
            f.close()
    except IOError, e:
        print >>sys.stderr, "Error writing output file `%s':" % outfile, e
        return
    if v(1):
        print "%s (%dx%d, %s bytes) -> %s (%dx%d, %s bytes)" % (infile, oldres[0], oldres[1], bytes(insize), outfile, newres[0], newres[1], bytes(len(data)))


################################################################################


def opterr(msg):
    print >>sys.stderr, "Error:", msg
    print >>sys.stderr, "Use `%s -h' to get help." % os.path.basename(sys.argv[0])
    sys.exit(2)

class OptionParser:
    def __init__(self, argument_callback=None):
        self.short_opts = {}
        self.long_opts = {}
        self.argument_callback = argument_callback

    def _parse_arg_name(self, name, arg_tail, int_tail):
        if name.endswith(int_tail):
            has_arg = True
            is_int = True
            name = name[:-len(int_tail)]
        elif name.endswith(arg_tail):
            has_arg = True
            is_int = False
            name = name[:-len(arg_tail)]
        else:
            has_arg = False
            is_int = False
        return (name, has_arg, is_int)

    def add_option(self, names, callback):
        for name in names.split(','):
            name = name.strip()
            if name.startswith("--"):
                name, has_arg, is_int = self._parse_arg_name(name[2:], "=", "=#")
                self.long_opts[name.lower()] = (callback, has_arg, is_int)
            elif name.startswith("-"):
                name, has_arg, is_int = self._parse_arg_name(name[1:], ":", "#")
                if len(name) != 1:
                    raise ValueError, "invalid short option name"
                self.short_opts[name] = (callback, has_arg, is_int)
            else:
                raise ValueError, "invalid option specifier"

    def get_opt(self, opt):
        try:
            if opt.startswith("--"):
                return self.long_opts[opt[2:].lower()]
            else:
                return self.short_opts[opt[1:]]
        except KeyError:
            opterr("unknown option `%s'" % opt)

    def handle_opt(self, opt, arg):
        cb, has_arg, is_int = self.get_opt(opt)
        if not has_arg:
            opterr("option `%s' requires no arguments" % opt)
        if is_int:
            try:
                arg = int(arg, 0)
            except ValueError:
                opterr("invalid integer value `%s' for option `%s'" % (arg, opt))
        cb(opt, arg)

    def run(self):
        args_only = False
        next_opt = None
        for arg in sys.argv[1:]:
            if args_only:
                self.argument_callback(arg)
            elif next_opt:
                self.handle_opt(next_opt, arg)
                next_opt = False
            elif arg == "--":
                args_only = True
            elif arg == "-":
                opterr("minus sign, but no option name")
            elif arg.startswith("--"):
                if '=' in arg:
                    self.handle_opt(*arg.split('=', 1))
                else:
                    cb, has_arg, is_int = self.get_opt(arg)
                    if has_arg:
                        next_opt = arg
                    else:
                        cb(arg)
            elif arg.startswith("-"):
                opt = arg[:2]
                arg = arg[2:]
                while True:
                    cb, has_arg, is_int = self.get_opt(opt)
                    if has_arg:
                        if arg:
                            self.handle_opt(opt, arg)
                        else:
                            next_opt = opt
                        break
                    else:
                        cb(opt)
                        if arg:
                            opt = "-" + arg[0]
                            arg = arg[1:]
                        else:
                            break
            else:
                self.argument_callback(arg)
        if next_opt:
            opterr("missing parameter for option `%s'" % next_opt)


def ocb_file(infile):
    global OutputFile, HadOutputFile
    if OutputFile:
        outfile = OutputFile
        HadOutputFile = True
    elif HadOutputFile:
        opterr("multiple arguments, but `--output' option has already been used")
    else:
        outfile = infile
    try:
        Process(infile, outfile)
    except KeyboardInterrupt:
        print >>sys.stderr, "Interrupt."
        sys.exit(1)

def ocb_res(opt, res):
    global MaxRes
    if res == '-':
        MaxRes = None
        print "res cleared"
        return
    res = res.lower()
    try:
        if 'x' in res:
            MaxRes = tuple(map(int, res.split('x', 1)))
        else:
            res = int(res)
            MaxRes = (res, res)
    except ValueError:
        opterr("invalid resolution `%s'" % res)
    if (MaxRes[0] < 1) or (MaxRes[1] < 1):
        opterr("invalid resolution `%s'" % res)
    if v(3): print "max resolution set to %dx%d" % MaxRes

def ocb_size(opt, size):
    global MaxSize, Quality
    scale = { 'k': 1024, 'm': 1024*1024 }.get(size[-1].lower(), 1)
    if scale > 1:
        size = size[:-1]
    try:
        MaxSize = int(scale * float(size))
    except ValueError:
        opterr("invalid file size `%s'" % size)
    Quality = None
    if v(3): print "site limit set to %d bytes" % MaxSize

def ocb_quality(opt, q):
    global MaxSize, Quality
    if (q < 1) or (q > 100):
        opterr("quality `%s' out of range" % q)
    MaxSize = None
    Quality = q
    if v(3): print "JPEG quality set to %d" % Quality

def ocb_output(opt, out):
    global OutputFile
    OutputFile = out

def ocb_pretend(dummy):
    global EnableWrite
    EnableWrite = False
    if v(3): print "not writing any files"

def ocb_verbose(dummy):
    global Verbosity
    Verbosity += 1
    if v(3): print "verbosity level set to", Verbosity

def ocb_force(dummy):
    global ForceRepack
    ForceRepack = True
    if v(3): print "forcing repack for all files"

def ocb_help(dummy):
    print __doc__
    print
    print USAGE
    print __help__.rstrip()
    examples()
    sys.exit(0)

def examples():
    print """
Examples:
  %s -r 640 -q 80 *.jpg
    -> re-encode all JPEG images to maximum 640x640 resolution
  %s -s 100k -o small.jpg big.jpg
    -> re-encode big.jpg into small.jpg so it's just below 100k in size
""" % (me, me)

if __name__ == "__main__":
    me = os.path.basename(sys.argv[0])
    USAGE = "Usage: %s [<options>] <image> [[<options>] <image> ...]" % me
    if len(sys.argv) < 2:
        print >>sys.stderr, USAGE
        examples()
        print >>sys.stderr, "Run `%s -h' to get full help." % me
        sys.exit(2)

    p = OptionParser(ocb_file)
    p.add_option("-h, --help", ocb_help)
    p.add_option("-r:, --res=", ocb_res)
    p.add_option("-s:, --size=", ocb_size)
    p.add_option("-q#, --quality=#", ocb_quality)
    p.add_option("-o:, --output=", ocb_output)
    p.add_option("-f, --force", ocb_force)
    p.add_option("-p, --pretend", ocb_pretend)
    p.add_option("-v, --verbose", ocb_verbose)
    p.run()