#!/usr/bin/python2

from __future__ import division

import os
import re
import shutil
import subprocess
import sys

TMP_DIR = None
DEST_DIR = None
SOURCE_DIR = None
DRY_RUN = False # this will still delete caches

VIDEO_FPS = 25
AUDIO_SAMPLE_RATE = 48000
ASPECT_RATIO = "4/3"


def mkarg(arg):
    if re.match("^[a-zA-Z0-9\-\\.,/@_:=]*$", arg):
        return arg

    if "'" not in arg:
        return "'%s'" % arg
    out = "\""
    for c in arg:
        if c in "\\$\"`":
            out += "\\"
        out += c
    out += "\""
    return out

def fail(line_count, msg):
    raise Exception(msg + " on line %d" % line_count)

def convert_frame_to_sample(frame):
    return frame * AUDIO_SAMPLE_RATE / VIDEO_FPS

def run_cmd(cmd):
    print "$", " ".join(map(mkarg, cmd))
    if DRY_RUN:
        return
    print
    ret = subprocess.Popen(cmd).wait()
    if ret != 0:
        print >>sys.stderr, "Failed on command", cmd
        raise Exception("Command returned non-zero: " + str(ret))

def read_file_contents(filename):
    try:
        f = open(filename)
        data = f.read().strip()
        f.close()
        return data
    except IOError:
        return None

def explode_video_to_png(source):
    image_cache_desc = os.path.join(TMP_DIR, "image_cache.txt")
    image_cache_dir = os.path.join(TMP_DIR, "image_cache")

    # Do nothing if the current cache is what we need
    current_cache = read_file_contents(image_cache_desc)
    if source == current_cache:
        return image_cache_dir

    # Remove if necessary
    if os.path.exists(image_cache_dir):
###     print "Confirm removal of image cache:", current_cache
###     ok = raw_input("(Y/n) ")
###     if ok != "Y":
###         print "Exiting..."
###         sys.exit(2)
        shutil.rmtree(image_cache_dir)

    cmd = [
        "mplayer",
        "-vo", "png:outdir=%s" % image_cache_dir,
        "-nosound",
        "-noconsolecontrols",
        "-noconfig", "user",
        "-benchmark",
        source,
    ]
    run_cmd(cmd)

    # Cache has been created, save the description
    f = open(image_cache_desc, "w")
    f.write(source)
    f.close()

    return image_cache_dir


def explode_video_to_wav(source):
    audio_cache_desc = os.path.join(TMP_DIR, "audio_cache.txt")
    audio_cache_file = os.path.join(TMP_DIR, "audio_cache.wav")

    # Do nothing if the current cache is what we need
    if source == read_file_contents(audio_cache_desc):
        return audio_cache_file

    cmd = [
        "mencoder",
        "-oac", "pcm",
        "-ovc", "copy",
        "-of", "rawaudio",
        "-o", audio_cache_file + ".raw",
        source,
    ]
    run_cmd(cmd)

    cmd = [
        "sox",
        "-r", str(AUDIO_SAMPLE_RATE), "-b", "16", "-e", "signed-integer",
        audio_cache_file + ".raw",
        audio_cache_file,
    ]
    run_cmd(cmd)

    # Cache has been created, save the description
    f = open(audio_cache_desc, "w")
    f.write(source)
    f.close()

    return audio_cache_file

def apply_audio_effects(source, dest, crop_start, crop_end, audio_normalize):
    cmd = [
        "sox",
        source,
        dest,
    ]
    if audio_normalize:
        cmd += ["gain", "-n"]
    if crop_start and crop_end:
        c = convert_frame_to_sample
        cmd += ["trim", "%ds" % c(crop_start), "%ds" % c(crop_end - crop_start)]
    run_cmd(cmd)

def apply_single_image_effects(source_file, dest_file, color_matrix):
    cmd = [
        "convert",
        source_file,
        "-color-matrix", color_matrix,
        dest_file,
    ]
    run_cmd(cmd)

def apply_image_effects(source_dir, crop_start, crop_end, color_matrix):
    dest_dir = os.path.join(TMP_DIR, "image_processed")
    if os.path.exists(dest_dir):
        shutil.rmtree(dest_dir)
    os.mkdir(dest_dir)

    inframe = crop_start
    outframe = 0
    while inframe <= crop_end:
        source_file = os.path.join(source_dir, str(inframe+1).zfill(8) + ".png")
        dest_file = os.path.join(dest_dir, str(outframe+1).zfill(8) + ".png")
        if color_matrix:
            apply_single_image_effects(source_file, dest_file, color_matrix)
        else:
            os.link(source_file, dest_file)
        inframe += 1
        outframe += 1

    return dest_dir

def combine_audio_video(audio_file, image_dir, dest):
    cmd = [
        "mencoder",
        "mf://%s/*.png" % image_dir,
        "-audiofile", audio_file,
        "-force-avi-aspect", ASPECT_RATIO,
        "-vf", "harddup",
        "-af", "channels=1",
        "-ovc", "lavc",
        "-lavcopts", "vcodec=ffv1:ilme:ildct",
        "-oac", "pcm",
        "-o", dest,
    ]
    run_cmd(cmd)


class Job(object):
    def __init__(self):
        self.source = None
        self.dest = None
        self.crop_start = None
        self.crop_end = None
        self.color_matrix = None
        self.audio_normalize = True

    def set_source(self, arg):
        self.source = os.path.join(SOURCE_DIR, arg)

    def set_dest(self, arg):
        self.dest = os.path.join(DEST_DIR, arg)
        if not self.dest.endswith(".avi"):
            self.dest += ".avi"

    def set_crop(self, arg):
        a, b = arg.split("-")
        self.crop_start = int(a)
        self.crop_end = int(b)

    def set_colormatrix(self, arg):
        [float(x) for x in arg.split(" ") if x] # check it's valid
        self.color_matrix = arg

    def set_whitecolor(self, arg):
        arg = arg.split(" ")
        color = arg[0]
        r = 0xff / int(color[0:2], 16)
        g = 0xff / int(color[2:4], 16)
        b = 0xff / int(color[4:6], 16)
        # don't change the brightness
        avg = (r + g + b) / 3
        if (avg - 1) > 0.02:
            diff = avg - 1.0
            r -= diff
            g -= diff
            b -= diff
        if len(arg) == 2:
            brightness = float(arg[1])
            r *= brightness
            g *= brightness
            b *= brightness
        self.set_colormatrix("%.3f 0 0  0 %.3f 0  0 0 %.3f" % (r, g, b))

    def set_audionormalize(self, arg):
        self.audio_normalize = int(arg)

    def validate(self, line_count, unique):
        if self.dest in unique:
            fail(line_count, "Non-unique output file: " + self.dest)
        if self.source is None:
            fail(line_count, "Missing source")
        if self.dest is None:
            fail(line_count, "Missing dest")
        if not os.path.isfile(self.source):
            fail(line_count, "Unable to find source: " + self.source)

    def is_done(self):
        return os.path.isfile(self.dest)

    def run(self):
        image_cache_dir = explode_video_to_png(self.source)
        image_dir = apply_image_effects(image_cache_dir, self.crop_start, self.crop_end, self.color_matrix)

        audio_cache_file = explode_video_to_wav(self.source)
        audio_file = os.path.join(TMP_DIR, "audio_processed.wav")
        apply_audio_effects(audio_cache_file, audio_file, self.crop_start, self.crop_end, self.audio_normalize)

        combine_audio_video(audio_file, image_dir, self.dest+".tmp")
        os.rename(self.dest+".tmp", self.dest)

    def __str__(self):
        return "Job :: %s (%s)" % (self.dest, self.source)

def main(frames):
    jobs = []
    unique = set()

    f = open(frames)

    job = None
    count = 0

    def append_job():
        if job is None:
            return
        job.validate(count, unique)
        if job.is_done():
            print "Skipping", job
        else:
            jobs.append(job)

    for line in f:
        count += 1
        line = line.strip()
        if line.startswith("#"):
            continue
        if not line:
            if job is not None:
                append_job()
                job = None
            continue

        if job is None:
            job = Job()
        cmd, arg = line.split(" ", 1)
        f = getattr(job, "set_"+cmd, None)
        if not f:
            fail(count, "Invalid command: " + cmd)
        try:
            f(arg)
        except Exception, e:
            fail(count, str(e))

    # trailing job...
    append_job()

    # optimise image and audio cache usage, use the current cache first if it exists
    current_image_cache = read_file_contents(os.path.join(TMP_DIR, "image_cache.txt"))
    jobs.sort(key=lambda job: (job.source if job.source != current_image_cache else "", job.dest))
    for job in jobs:
        print "\n\n\nStarted job:", job, "\n\n"
        job.run()

if __name__ == "__main__":
    try:
        frames = sys.argv[1]
        SOURCE_DIR = sys.argv[2]
        DEST_DIR = sys.argv[3]
        TMP_DIR = sys.argv[4]
    except IndexError:
        print >>sys.stderr, "Usage: %s frames.txt source_dir dest_dir tmp_dir" % sys.argv[0]
        sys.exit(1)

    main(frames)