#!/usr/bin/python
# coding=utf-8

import math
import os, os.path
import re
import sys
import time


exclude = []
#exclude = ['foo', 'bar'] # Don't convert files whose title contains one of these strings
ffmpeg_bin = 'ffmpeg' # Path to ffmpeg binary
projectx_jar = '/usr/src/ProjectX_Source_0.90.4/ProjectX.jar' # Path to ProjectX.jar
max_filename_length = 62 # The Viewty doesn't show files with file names > 62 characters
vcodec = 'xvid' # Video codec
vrate = 700 # Video bitrate
vtwopass = False # Do two-pass encoding?
vtag = 'DX50' # Video tag override (Fake as DivX 5)
acodec = 'mp3' # Audio codec
arate = 112 # Audio bitrate
asampling = 44100 # Audio sample rate


def get_recording_info(directory):
    """
    Retrieves some information about a recording.
    Returns a dict with the keys `title`, `subtitle`, `description`, `datetime`.
    The date is the date and time of the recording in the format `YYYY-MM-DD hh:mm`.
    Note, that subtitle and description may be missing.
    """
    result = {}
    info = open(directory + 'info.vdr')
    if info:
        line = info.readline()
        while line:
            (code, content) = line.split(' ', 1)
            if code == 'E':
                result['datetime'] = time.strftime(r'%Y-%m-%d %H:%M', time.localtime(int(content.split(' ')[1])))
            elif code == 'T':
                result['title'] = content
            elif code == 'S':
                result['subtitle'] = content
            elif code == 'D':
                result['description'] = content
            line = info.readline()
    return result


def demux(directory):
    """
    Demultiplex the VDR files in a specified directory.
    Returns the exit code of the process.
    """
    # How many raw files exist?
    files = []
    i = 1
    while os.path.exists(directory + '%03i.vdr' % (i)):
        files.append(directory + '%03i.vdr' % (i))
        i +=1
    print "Files: " + ", ".join(files)
    # Run demuxer
    demux_cmd = ['java', '-jar', projectx_jar, '-name', 'demuxed', '-out', directory, '-demux'] + files
    return os.spawnvp(os.P_WAIT, demux_cmd[0], demux_cmd)

def extract_aspect_from_log(directory):
    """
    Extracts and returns the aspect ratio from a project x log file.
    """
    pattern = re.compile(r'video basics.*\((\d+):(\d+)\)')
    aspect = None
    
    log = open(directory + 'demuxed_log.txt', 'r')
    buf = log.readline()
    while buf:
        match = pattern.search(buf)
        if match:
            return float(match.group(1)) / float(match.group(2))
        buf = log.readline()
    log.close()


def calculate_maximum_size(aspect, max_pixels):
    """
    Calculates the maxium size for a given aspect ratio and a constaint with
    a maximum number of total pixels.
    """
    height = int(math.sqrt(max_pixels / aspect))
    width = int(aspect * height)
    # Ensure even frame sizes
    if height % 2:
        height -= 1
    if width % 2:
       width -= 1
    return str(width) + 'x' + str(height)


def encode(directory, size):
    """
    Encodes the demuxed files using FFmpeg and the parameters specified at the
    top of this script.
    Expects the directory with the demuxed files and the target size in the
    format `widthxheight` (e.g. `640x480`).
    Returns the exit code of the process.
    """
    info = get_recording_info(directory)
    filename = info.get('title', '') + ' ' + info['datetime'] + ' ' + info.get('subtitle', '')
    filename = filename[:max_filename_length-4] + '.avi'
    # Escape target filename
    pattern = re.compile(r'[^\w\.\,\-]+')
    filename = pattern.sub('_', filename)
    print "Target filename:", filename
    # Build general encoding command
    encode_cmd = [
        ffmpeg_bin, '-i', directory + 'demuxed.m2v', '-i', directory + 'demuxed.mp2',
        '-vcodec', vcodec, '-s', size, '-b', str(vrate) + 'kb', '-vtag', vtag,
        '-acodec', acodec, '-ab', str(arate) + 'kb', '-ar', str(asampling),
        '-y', directory + filename
    ]
    # Run FFmpeg. Arguments must all be strings!
    if vtwopass:
        # Two-pass encoding
        encode_twopass_cmd = encode_cmd[:1] + ['-passlogfile',  directory + 'pass.log', '-pass', '1'] + encode_cmd[1:]
        status = os.spawnvp(os.P_WAIT, encode_twopass_cmd[0], encode_twopass_cmd)
        if status != 0:
            return status
        encode_twopass_cmd[4] = '2'
        return os.spawnvp(os.P_WAIT, encode_twopass_cmd[0], encode_twopass_cmd)
    else:
        return os.spawnvp(os.P_WAIT, encode_cmd[0], encode_cmd)


def cleanup(directory):
    """
    Deletes temporary files.
    """
    try:
        os.remove(directory + 'demuxed.m2v')
        os.remove(directory + 'demuxed.mp2')
        os.remove(directory + 'demuxed_log.txt')
        os.remove(directory + 'pass.log-0.log')
    except OSError, e:
        pass


def main():
    # Only run after recording has finished
    if len(sys.argv) < 3 or sys.argv[1] != 'after':
        print "Only running after recording. Exiting."
        sys.exit(0)
    
    directory = sys.argv[2]
    if not directory.endswith('/'):
        directory += '/'
    
    # Get info
    info = get_recording_info(directory)
    
    # Exclude? Each exclusion entry must be encoded to ISO-8859-1
    if filter(lambda x: unicode(x, 'utf-8').encode('latin-1').lower() in info.get('title', '').lower(), exclude):
        print >>sys.stderr, "Not processing this recording. Matched an exclusion filter."
        sys.exit(0)
    
    print "Demuxing..."
    status = demux(directory)
    if status != 0:
        print >>sys.stderr, "Error while demuxing. Status code %i" % (status)
        sys.exit(1)
    
    aspect = extract_aspect_from_log(directory)
    if not aspect:
        print >>sys.stderr, "Could not determine aspect ratio from demux output!"
        sys.exit(1)
    print "Aspect:", aspect
    
    size = calculate_maximum_size(aspect, 640 * 480)
    print "Target size:", size
    
    print "Encoding..."
    status = encode(directory, size)
    if status != 0:
        print >>sys.stderr, "Error while encoding. Status code %i" % (status)
        sys.exit(1)
    
    # Clean up
    print "Cleaning up..."
    cleanup(directory)
    
    print "Done."


if __name__ == "__main__":
    main()