#!/usr/local/bin/bash
#
# Copyright (c) 2008  SlashOrg.Net <dev@slashorg.net>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#
# Filename mkv2mp4.sh
# Version 1.03
# Last Update 9 dec 2008
# Converts mkv files with h264 video to mp4 for PS3 without reencoding.
# Send bug reports to dev@slashorg.net
# 
# * Version history *
#  - 1.00
#  - 1.01 : added support to update H264 level for PS3 compatibility.
#  - 1.02 : handle multiple audio streams in mkv file.
#  - 1.03 : added mp4creator support

# mp4box or mp4creator.
MP4MUX=mp4creator
MAXSIZE=4500000000
#MAXSIZE=1000000000

# Nothing much to modify here.
PATH=/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin
SRCFILE=$1
DESTDIR=$2
TMPDIR=`dirname "$SRCFILE"`
TMPFILEBASE="$TMPDIR/`date +%H%M%S`$$"
MAGIC_H264=67640029

SEARCH=`echo "$TMPFILEBASE" | sed 's/.*\///'`
COPYTO=`echo "$SRCFILE" | sed 's/.*\///'`


# Functions
function checkfor {
	which $1 1>/dev/null 2>&1

	if [ $? -ne 0 ]; then
		echo "$1 cannot be found."
		echo "Quitting."
		exit
	fi
}

function dispmsg {
	MESSAGE="$1"

	echo
	echo "****************************************"
	echo " *** $MESSAGE ***"
	echo "****************************************"
	echo
}

function cleanup {
	SRC="$1"

	echo "Deleting temp files..."

	rm -f "$SRC"
	rm -f "$SRC.h264"
	rm -f "$SRC.aud"
	rm -f "$SRC.aac"
}

function stoponerror {
	RETCODE=$1
	MESSAGE=$2

	if [ $RETCODE -ne 0 ]; then
		if [ "$MESSAGE" ]; then
			echo "An error occured in $MESSAGE."
		else
			echo "An error occured."
		fi
		echo "Quitting."
		exit
	fi
}

# Check input file
if [ "$SRCFILE" == "" ]; then
	echo "Usage : mkv2mp4.sh /path/to/video.mkv [/copy/to/folder]"
	exit
fi

FILECHK=`file $SRCFILE 2>/dev/null | grep Matroska`

if [ "$FILECHK" == "" ]; then
	dispmsg "$SRCFILE is not mkv file or does not exist"
	exit
fi

if [[ "$DESTDIR" != "" && ! -d "$DESTDIR" ]]; then
	dispmsg "$DESTDIR is not a folder."
	echo -n "Create ? (y/n) "
	read ANS
	if [ "$ANS" == "y" ]; then
		mkdir -p "$DESTDIR"
		stoponerror $? "making new folder"
	else
		echo "Quitting."
		exit
	fi
fi

# Check for programs.
checkfor mkvmerge
checkfor mkvextract
checkfor mkvinfo
checkfor ffmpeg
checkfor hexdump

if [ "$MP4MUX" == "mp4creator" ]; then
	checkfor mp4creator
else
	checkfor mp4box
fi

# Find audio/video track IDs and FPS (used by mp4box).
VIDID=`mkvmerge -i "$SRCFILE" | grep V_MPEG4 | sed -e 's/Track ID //' -e 's/:.*//'`
AUDID=`mkvmerge -i "$SRCFILE" | grep -E 'A_DTS|A_AC3' | sed -e 's/Track ID //' -e 's/:.*//'`
FPS=`mkvinfo "$SRCFILE" | grep fps | head -n 1 | sed -e 's/.*(//' -e 's/ fps.*//'`

if [[ "$VIDID" == "" || "$AUDID" == "" ]]; then
	echo "Ouuupps. Could not get track IDs."
	echo "Quitting."
	exit
fi

if [ `echo "$AUDID" | wc -l` -ne 1 ]; then
	dispmsg "Found multiple audio streams."

	echo `mkvmerge -i "$SRCFILE" | grep -E 'A_DTS|A_AC3'`

	OK=0
	while [ $OK -ne 1 ]; do
		echo -n "Choose a stream ID : "
		read ANS

		if [ "`echo $ANS | sed 's/[^0-9]//g'`" != "" ]; then
			if [ "`mkvmerge -i \"$SRCFILE\" | grep -E 'A_DTS|A_AC3' | sed -e 's/Track ID //' -e 's/:.*//' | grep $ANS`" != "" ]; then
				OK=1
				AUDID=$ANS
			else
				echo "  $ANS is not a valid audio ID."
			fi
		fi
	done
	
fi

dispmsg "Found track IDs : Video ID $VIDID ; Audio ID : $AUDID"

# Check file size.
if [ `ls -l "$SRCFILE" | awk '{ print $5; }'` -gt $MAXSIZE ]; then
	dispmsg "Warning. Splitting will occur."

	mkvmerge -o "$TMPFILEBASE.mkv" --split size:$MAXSIZE "$SRCFILE"
	stoponerror $? "mkvmerge"
else
	ln -s "$COPYTO" "$TMPFILEBASE.mkv"
fi


I=1
for ONEFILE in `ls "$TMPDIR" | grep $SEARCH | grep mkv$`; do
	dispmsg "Processing batch $I..."

	# Extract audio and video.
	dispmsg "Extracting audio and video..."
	mkvextract tracks "$TMPDIR/$ONEFILE" $VIDID:"$TMPDIR/$ONEFILE.h264" $AUDID:"$TMPDIR/$ONEFILE.aud"
	stoponerror $? "mkvextract"

	H264HEX=`hexdump -s 4 -n 4 -x "$TMPDIR/$ONEFILE.h264" | head -n 1 | awk '{ print $2$3; }' | sed -E 's/([0-9]{2})([0-9]{2})([0-9]{2})([0-9]{2})/\2\1\4\3/g'`

	if [ "$H264HEX" != "$MAGIC_H264" ]; then
		dispmsg "Making h264 compatible with PS3..."
		printf "`echo -n "$MAGIC_H264" | sed -E 's/([0-9]{2})([0-9]{2})([0-9]{2})([0-9]{2})/\\\x\1\\\x\2\\\x\3\\\x\4/g'`" | dd of="$TMPDIR/$ONEFILE.h264" bs=1 seek=4 count=4 conv=notrunc
	fi

	# Recompress audio.
	dispmsg "Recompressing audio to AAC..."
	ffmpeg -i "$TMPDIR/$ONEFILE.aud" -ac 2 -ab 192k "$TMPDIR/$ONEFILE.aac"
	stoponerror $? "FFMPEG"

	# Mux to MP4.
	dispmsg "Building MP4 file using $MP4MUX..."

	if [ "$MP4MUX" == "mp4creator" ]; then
		echo "Adding video..."
		mp4creator -create="$TMPDIR/$ONEFILE.h264" -rate=$FPS "$TMPDIR/$ONEFILE.m4v"
		stoponerror $? "$MP4MUX"

		echo "Adding audio..."
		mp4creator -create="$TMPDIR/$ONEFILE.aac" -interleave "$TMPDIR/$ONEFILE.m4v"
	else
		mp4box -new "$TMPDIR/$ONEFILE.m4v" -add "$TMPDIR/$ONEFILE.h264" -add "$TMPDIR/$ONEFILE.aac" -tmp "$TMPDIR" -fps $FPS
	fi
	stoponerror $? "$MP4MUX"

	# Bug in mp4box with long filenames...

	if [ "$DESTDIR" == "" ]; then
		dispmsg "Copying file to $TMPDIR/$COPYTO-$I.m4v"
		mv "$TMPDIR/$ONEFILE.m4v" "$TMPDIR/$COPYTO-$I.m4v" 2>/dev/null
	else
		dispmsg "Copying file to $DESTDIR/$COPYTO-$I.m4v"
		mv "$TMPDIR/$ONEFILE.m4v" "$DESTDIR/$COPYTO-$I.m4v" 2>/dev/null
	fi
	stoponerror $? "moving m4v files"

	# Delete temp files.
	cleanup "$TMPDIR/$ONEFILE"
	I=$(( I + 1 ))
done


dispmsg "All done."

exit

