#!/bin/bash
# mkvdts2ac3.sh - add an AC3 track to an MKV from its DTS track
# Author: Jake Wharton <jakewharton@gmail.com>
# Chris Hoekstra <chris.hoekstra@gmail.com>
# Website: http://jakewharton.com
# http://github.com/JakeWharton/mkvdts2ac3/
# Version: 1.5.1
# License:
# Copyright 2010 Jake Wharton
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#Prerequisites
# http://www.networkedmediatank.com/wiki/index.php/Converting_DTS_to_AC3
#Download and install the following:

    #mkvtoolnix - Matroska (MKV) tools
    #libdca - DTS to WAV decoder
    #aften - WAV to AC3 encoder
    #libdca-utils


# Display version header
echo "mkvdts2ac3-1.5.1 - by Jake Wharton <jakewharton@gmail.com> and"
echo " Chris Hoekstra <chris.hoekstra@gmail.com>"
echo

# Debugging flags
# DO NOT EDIT THESE! USE --debug OR --test ARGUMENT INSTEAD.
PRINT=0
PAUSE=0
EXECUTE=1

# Default values
PRIORITY=0
FORCE=0
NOCOLOR=0
MD5=0
INITIAL=0
NEW=0
WD="/tmp" # Working Directory (Use the -w/--wd argument to change)
DUCMD="$(which \du) -k"

#---------- FUNCTIONS --------
displayhelp() {
# Usage: displayhelp
echo "Usage: `basename $0` [options] <filename>"
echo "Options:"
echo " -c TITLE, Custom AC3 track title."
echo " --custom TITLE"
echo " -d, --default Mark AC3 track as default."
echo " -e, --external Leave AC3 track out of file. Does not modify the"
echo " original matroska file. This overrides '-n' and"
echo " '-d' arguments."
echo " -f, --force Force processing when AC3 track is detected"
echo " -i, --initial New AC3 track will be first in the file."
echo " -k, --keep-dts Keep external DTS track (implies '-n')."
echo " -m, --nocolor Do not use colors (monotone)."
echo " --md5 Perform MD5 comparison when copying across drives."
echo " -n, --no-dts Do not retain the DTS track."
echo " --new Do not copy over original. Create new adjacent file."
echo " -o MODE Pass a custom audio output mode to libdca."
echo " -p PRIORITY Modify niceness of executed commands."
echo " -t TRACKID,"
echo " --track TRACKID Specify alternate DTS track."
echo " -w FOLDER,"
echo " --wd FOLDER Specify alternate temporary working directory."
echo
echo " --test Print commands only, execute nothing."
echo " --debug Print commands and pause before executing each."
echo
echo " -h, --help Print command usage."
echo " -v, --verbose Turn on verbose output"
echo " -V, --version Print script version information."
}


# Usage: color shade
function color {
# Are we in Cron?
if [ ! -t 0 -o $NOCOLOR = 1 ]; then return 1; fi
case $1 in
off|OFF) echo -n '';;
red|RED) echo -n '';;
green|GREEN) echo -n '';;
yellow|YELLOW) echo -n '';;
bell|BELL) echo -n '';;
*) ;;
esac
}

# Usage: timestamp ["String to preface time display"]
timestamp() {
CURRTIME=$(date +%s)
secs=$(($CURRTIME - $PREVTIME))
PREVTIME=$((CURRTIME))
h=$(( secs / 3600 ))
m=$(( ( secs / 60 ) % 60 ))
s=$(( secs % 60 ))

if [ $EXECUTE = 1 -a $PAUSE = 0 ]; then
echo -n "$1 "
printf "%02d:%02d:%02d " $h $m $s
echo "($secs seconds)\n"
fi
}

# Usage: error "String to display"
error() {
color BELL;color RED
echo "$1"
color OFF
}

# Usage: dopause
dopause() {
if [ $PAUSE = 1 ]; then
read
fi
}

# Usage: checkdep appname
checkdep() {
if [ -z "$(which $1)" -o ! -x "$(which $1)" ]; then
error "ERROR: The program '$1' is not in the path. Is $1 installed?"
exit 1
fi
}

# Usage: cleanup file
cleanup() {
if [ -f $1 ]; then
rm $1
if [ $? -ne 0 ]; then
"There was a problem removing the file \"$1\". Please remove manually."
return 1
fi
fi
}

# Usage: checkerror $? "Error message to display" [exit on error:BOOL]
checkerror() {
if [ $1 -ne 0 ]; then
error "$2"
# if optional BOOL then exit, otherwise return errorcode 1
if [ $3 -gt 0 ]; then
# honor KEEPDTS
if [ -z $KEEPDTS ]; then
cleanup $DTSFILE
fi

cleanup $AC3FILE
cleanup $TCFILE
exit 1
fi
return 1
fi
}

# Usage: doprint "String to print"
doprint() {
if [ $PRINT = 1 ]; then
echo -e "$1"
fi
}

#---------- START OF PROGRAM ----------
# Start the timer and make a working copy for future timings
START=$(date +%s)
PREVTIME=$(($START))

# Parse arguments and/or filename
while [ -z "$MKVFILE" ]; do

# If we're out of arguments no filename was passed
if [ $# -eq 0 ]; then
error "ERROR: You must supply a filename."
echo ""
displayhelp
exit 1
fi

case "$1" in
"-c" | "--custom" ) # Use custom name for AC3 track
shift
DTSNAME=$1
;;
"-d" | "--default" ) # Only allow this if we aren't making the file external
if [ -z $EXTERNAL ]; then
DEFAULT=1
fi
;;
"-e" | "--external" ) # Don't allow -d or -n switches if they're already set
EXTERNAL=1
NODTS=0
KEEPDTS=0
DEFAULT=0
;;
"-f" | "--force" ) # Test for AC3 track exits immediately. Use this to continue
FORCE=1
;;
"-i" | "--initial" ) # Make new AC3 track the first in the file
INITIAL=1
;;
"-k" | "--keep-dts" ) # Only allow external DTS track if muxing AC3 track
if [ -z $EXTERNAL ]; then
KEEPDTS=1
fi
;;
"-m" | "--nocolor" | "--monotone" ) # Turns off colors
NOCOLOR=1
;;
"--md5" ) #Perform MD5 comparison when copying across drives
MD5=1
;;
"-n" | "--no-dts" ) # Only allow this if we aren't making the file external
if [ -z $EXTERNAL ]; then
NODTS=1
fi
;;
"--new" ) # Do not overwrite original. Create new adjacent file.
NEW=1
;;
"-o" ) # Move required audio mode value "up"
shift
AUDIOMODE=$1
;;
"-p" ) # Move required priority value "up"
shift
PRIORITY=$1
;;
"-t" | "--track" ) # Move required TRACKID argument "up"
shift
DTSTRACK=$1
;;
"-w" | "--wd" ) # Specify working directory manually
shift
WD=$1
;;
"--test" ) # Echo commands and do not execute
if [ $PAUSE = 1 ]; then
error "WARNING: --test overrides previous --debug flag."
fi
PRINT=1
EXECUTE=0
;;
"--debug" ) # Echo commands and pause before executing
if [ $EXECUTE = 0 ]; then
error "ERROR: --debug flag not valid with --test."
displayhelp
exit 1
fi
PRINT=1
PAUSE=1
EXECUTE=1
;;
"-h" | "--help" )
displayhelp
exit 0
;;
"-v" | "--verbose" ) # Turn on verbosity
PRINT=1
;;
"-V" | "--version" ) # Version information is always displayed so just exit here
exit 0
;;
-* | --* )
error "ERROR: Invalid argument '$1'."
echo ""
displayhelp
exit 1
;;
* )
MKVFILE=$1
shift

# Ensure there are no arguments after the filename
if [ $# -ne 0 ]; then
error "ERROR: You cannot supply any arguments after the filename. Please check the command syntax below against what has been parsed."
echo ""
echo "Control Flags:"
echo " Strip DTS: $NODTS"
echo " Keep DTS: $KEEPDTS"
echo " Set AC3 default: $DEFAULT"
echo " External AC3: $EXTERNAL"
echo " DTS track: $DTSTRACK"
echo " MKV file: $MKVFILE"
echo ""
echo "Debugging Flags:"
echo " Print commands: $PRINT"
echo " Pause after print: $PAUSE"
echo " Execute commands: $EXECUTE"
echo ""
displayhelp
exit 1
fi
;;
esac

# Move arguments "up" one spot
shift
done

# File and dependency checks
if [ $EXECUTE = 1 ]; then
# Check the file exists and we have permissions
if [ ! -f "$MKVFILE" ]; then
error "ERROR: '$MKVFILE' is not a file."
exit 1
elif [ ! -r "$MKVFILE" ]; then
error "ERROR: Cannot read '$MKVFILE'."
exit 1
elif [ -z $EXTERNAL ]; then
if [ ! -w "$MKVFILE" ]; then
# Only check write permission if we're not keeping the AC3 external
error "ERROR: Cannot write '$MKVFILE'."
exit 1
fi
fi

# Check dependencies
checkdep mkvmerge
checkdep mkvextract
checkdep mkvinfo
checkdep dcadec
checkdep aften
checkdep rsync
fi

# Added check to see if AC3 track exists. If so, no need to continue
if [ "$(mkvmerge -i "$MKVFILE" | grep -i "A_AC3")" ]; then
echo "AC3 track already exists in $MKVFILE."
if [ $FORCE = 0 ]; then
echo "Use -f or --force argument to bypass this check."
exit 1
fi
echo "Force mode is on. Continuing..."
fi

# Path to file
DEST=$(dirname "$MKVFILE")

# File name without the extension
NAME=$(basename "$MKVFILE" .mkv)

# Setup temporary files
DTSFILE="$WD/$NAME.dts"
AC3FILE="$WD/$NAME.ac3"
TCFILE="$WD/$NAME.tc"
NEWFILE="$WD/$NAME.new.mkv"

doprint "MKV FILE: $MKVFILE
DTS FILE: $DTSFILE
AC3 FILE: $AC3FILE
TIMECODE: $TCFILE
NEW FILE: $NEWFILE
WORKING DIRECTORY: $WD"

# ------ GATHER DATA ------
# If the track id wasn't specified via command line then search for the first DTS audio track
if [ -z $DTSTRACK ]; then
doprint "\nFind first DTS track in MKV file."
doprint "> mkvmerge -i \"$MKVFILE\" | grep -m 1 \"audio (A_DTS)\" | cut -d ":" -f 1 | cut -d \" \" -f 3"
DTSTRACK="DTSTRACK" #Value for debugging
dopause
if [ $EXECUTE = 1 ]; then
DTSTRACK=$(mkvmerge -i "$MKVFILE" | grep -m 1 "audio (A_DTS)" | cut -d ":" -f 1 | cut -d " " -f 3)

# Check to make sure there is a DTS track in the MVK
if [ -z $DTSTRACK ]; then
error "ERROR: There are no DTS tracks in '$MKVFILE'."
exit 1
fi
fi
else
# Checks to make sure the command line argument track id is valid
doprint "Checking to see if DTS track specified via arguments is valid."
doprint "> mkvmerge -i \"$MKVFILE\" | grep \"Track ID $DTSTRACK: audio (A_DTS)\""
dopause
if [ $EXECUTE = 1 ]; then
VALID=$(mkvmerge -i "$MKVFILE" | grep "Track ID $DTSTRACK: audio (A_DTS)")

if [ -z "$VALID" ]; then
error "ERROR: Track ID '$DTSTRACK' is not a DTS track and/or does not exist."
exit 1
else
echo "INFO: Using alternate DTS track with ID '$DTSTRACK'."
fi
fi
fi
# Get the specified DTS track's information
doprint "\nExtract track information for selected DTS track.
> mkvinfo \"$MKVFILE\" | grep -A 25 \"Track number: $DTSTRACK\""

INFO="INFO" #Value for debugging
dopause
if [ $EXECUTE = 1 ]; then
INFO=$(mkvinfo "$MKVFILE" | grep -A 25 "Track number: $DTSTRACK")
fi

#Get the language for the DTS track specified
doprint "\nExtract language from track info.
> echo \"$INFO\" | grep -m 1 \"Language\" | cut -d \" \" -f 5"

DTSLANG="DTSLANG" #Value for debugging
dopause
if [ $EXECUTE = 1 ]; then
DTSLANG=$(echo "$INFO" | grep -m 1 "Language" | cut -d " " -f 5)
fi

# Check if a custom name was already specified
if [ -z $DTSNAME ]; then
# Get the name for the DTS track specified
doprint "\nExtract name for selected DTS track. Change DTS to AC3 and update bitrate if present."
doprint "> echo \"$INFO\" | grep -m 1 \"Name\" | cut -d \" \" -f 5- | sed \"s/DTS/AC3/\" | awk '{gsub(/[0-9]+(\.[0-9]+)?(M|K)bps/,\"448Kbps\")}1'"
DTSNAME="DTSNAME" #Value for debugging
dopause
if [ $EXECUTE = 1 ]; then
DTSNAME=$(echo "$INFO" | grep -m 1 "Name" | cut -d " " -f 5- | sed "s/DTS/AC3/" | awk '{gsub(/[0-9]+(\.[0-9]+)?(M|K)bps/,"448Kbps")}1')
fi
fi

# ------ EXTRACTION ------
# Extract timecode information for the target track
doprint "\nExtract timecode information for the audio track.
> mkvextract timecodes_v2 \"$MKVFILE\" $DTSTRACK:\"$TCFILE\"
> sed -n \"2p\" \"$TCFILE\"
> rm -f \"$TCFILE\""

DELAY="DELAY" #Value for debugging
dopause
if [ $EXECUTE = 1 ]; then
color YELLOW; echo "Extracting Timecodes:"; color OFF
nice -n $PRIORITY mkvextract timecodes_v2 "$MKVFILE" $DTSTRACK:"$TCFILE"
DELAY=$(sed -n "2p" "$TCFILE")
cleanup $TCFILE
timestamp "Extract timecodes took: "
fi

# Extract the DTS track
doprint "\nExtract DTS file from MKV.
> mkvextract tracks \"$MKVFILE\" $DTSTRACK:\"$DTSFILE\""

dopause
if [ $EXECUTE = 1 ]; then
color YELLOW; echo "Extracting DTS Track: "; color OFF
nice -n $PRIORITY mkvextract tracks "$MKVFILE" $DTSTRACK:"$DTSFILE" |grep -v CodecID
checkerror $? "ERROR: Extracting DTS track failed." 1
timestamp "Extract DTS track took: "
fi

# ------ CONVERSION ------
# Convert DTS to AC3
doprint "Converting DTS to AC3.
> dcadec -o wavall \"$DTSFILE\" | aften - \"$AC3FILE\""

dopause
if [ $EXECUTE = 1 ]; then
if [ -z $AUDIOMODE ]; then
AUDIOMODE="wavall"
fi

color YELLOW; echo "Converting DTS to AC3:"; color OFF
nice -n $PRIORITY dcadec -o $AUDIOMODE "$DTSFILE" | nice -n $PRIORITY aften -v 0 - "$AC3FILE"
checkerror $? "ERROR: Converting the DTS file to AC3 failed" 1
DTSFILESIZE=$($DUCMD "$DTSFILE" | cut -f1) # Capture DTS filesize for end summary

# If we are keeping the DTS track external do not delete it
if [ -z $KEEPDTS ]; then
cleanup $DTSFILE
fi

timestamp "Convert DTS track took: "
fi

# Check there is enough free space for AC3+MKV
if [ $EXECUTE = 1 ]; then
MKVFILESIZE=$($DUCMD "$MKVFILE" | cut -f1)
AC3FILESIZE=$($DUCMD "$AC3FILE" | cut -f1)
WDFREESPACE=$(\df -B 1 "$WD" | tail -1 | awk '{print $4}')
if [ $(($MKVFILESIZE + $AC3FILESIZE)) -gt $WDFREESPACE ]; then
error "ERROR: There is not enough free space on \"$WD\" to create the new file."
exit 1
fi
fi

if [ $EXTERNAL ]; then
# We need to trick the rest of the script so that there isn't a lot of
# code duplication. Basically $NEWFILE will be the AC3 track and we'll
# change $MKVFILE to where we want the AC3 track to be so we don't
# overwrite the MKV file only an AC3 track
NEWFILE=$AC3FILE
MKVFILE="$DEST/$NAME.ac3"
else
# Start to "build" command
CMD="nice -n $PRIORITY mkvmerge -q "

# Puts the AC3 track as the second in the file if indicated as initial
if [ $INITIAL = 1 ]; then
CMD="$CMD --track-order 0:1,1:0"
fi

# Declare output file
CMD="$CMD -o \"$NEWFILE\""

# If user doesn't want the original DTS track drop it
if [ $NODTS ]; then
# Count the number of audio tracks in the file
AUDIOTRACKS=$(mkvmerge -i "$MKVFILE" | grep "audio (A_" | wc -l)

if [ $AUDIOTRACKS -eq 1 ]; then
# If there is only the DTS audio track then drop all audio tracks
CMD="$CMD -A"
else
# Get a list of all the other audio tracks
SAVETRACKS=$(mkvmerge -i "$MKVFILE" | grep "audio (A_" | cut -d ":" -f 1 | grep -vx "Track ID $DTSTRACK" | cut -d " " -f 3 | awk '{ if (T == "") T=$1; else T=T","$1 } END { print T }')
# And copy only those
CMD="$CMD -a \"$SAVETRACKS\""
fi
fi

# Add original MKV file to the MKV command
CMD="$CMD \"$MKVFILE\""


# If user wants new AC3 as default then add appropriate arguments to command
if [ $DEFAULT ]; then
CMD="$CMD --default-track 0"
fi

# If the language was set for the original DTS track set it for the AC3
if [ $DTSLANG ]; then
CMD="$CMD --language 0:$DTSLANG"
fi

# If the name was set for the original DTS track set it for the AC3
if [ "$DTSNAME" ]; then
CMD="$CMD --track-name 0:\"$DTSNAME\""
fi

# If there was a delay on the original DTS set the delay for the new AC3
if [ $DELAY != 0 ]; then
CMD="$CMD --sync 0:$DELAY"
fi

# Append new AC3
CMD="$CMD \"$AC3FILE\""

# ------ MUXING ------
# Run it!
doprint "\nRunning main remux."
doprint "> $CMD"
dopause
if [ $EXECUTE = 1 ]; then
color YELLOW; echo "Muxing AC3 Track in:"; color OFF
eval $CMD
checkerror $? "ERROR: Merging the AC3 track back into the MKV failed." 1
fi

# Delete AC3 file if successful
doprint "Removing temporary AC3 file."
doprint "> rm -f \"$AC3FILE\""
dopause
cleanup $AC3FILE
fi

# If we are creating an adjacent file adjust the name of the original
if [ $NEW = 1 ]; then
MKVFILE="$DEST/$NAME-AC3.mkv"
fi

# Check to see if the two files are on the same device
NEWFILEDEVICE=$(\df "$WD" | tail -1 | cut -d" " -f1)
DSTFILEDEVICE=$(\df "$DEST" | tail -1 | cut -d" " -f1)

if [ "$NEWFILEDEVICE" = "$DSTFILEDEVICE" ]; then
# If we're working on the same device just move the file over the old one
if [ "$NEWFILE" = "$MKVFILE" ]; then
doprint "\nNew file and destination are the same. No action is required."
else
doprint "\nMoving new file over old one."
doprint "> mv \"$NEWFILE\" \"$MKVFILE\""
dopause
if [ $EXECUTE = 1 ]; then
color YELLOW; echo "MOVING new file over old file. DO NOT KILL THIS PROCESS OR YOU WILL EXPERIENCE DATA LOSS!"; color OFF
echo "NEW FILE: $NEWFILE"
echo "MKV FILE: $MKVFILE"
mv "$NEWFILE" "$MKVFILE"
checkerror $? "ERROR: There was an error copying the new MKV over the old one. You can perform this manually by moving '$NEWFILE' over '$MKVFILE'."
fi
fi
else
doprint "\nCopying new file over the old one."
doprint "> cp \"$NEWFILE\" \"$MKVFILE\""
dopause

# Check there is enough free space for the new file
if [ $EXECUTE = 1 ]; then
MKVFILEDIFF=$(($($DUCMD "$NEWFILE"|cut -f1) - $MKVFILESIZE))
DESTFREESPACE=$(\df -B 1 "$DEST" | tail -1 | awk '{print $4}')
if [ $MKVFILEDIFF -gt $DESTFREESPACE ]; then
error "ERROR: There is not enough free space to copy the new MKV over the old one. Free up some space and then copy '$NEWFILE' over '$MKVFILE'."
exit 1
fi

# Rsync our new MKV with the AC3 over the old one OR if we're using the -e
# switch then this actually copies the AC3 file to the original directory
color YELLOW; echo "Moving new file over old file. DO NOT KILL THIS PROCESS OR YOU WILL EXPERIENCE DATA LOSS!"; color OFF
rsync -av "$NEWFILE" "$MKVFILE"
checkerror $? "ERROR: There was an error copying the new MKV over the old one. You can perform this manually by copying '$NEWFILE' over '$MKVFILE'." 1

if [ $MD5 = 1 ]; then
# Check MD5s are equal to ensure the full file was copied (because du sucks across filesystems and platforms)
OLDFILEMD5=$(md5sum "$NEWFILE" | cut -d" " -f1)
NEWFILEMD5=$(md5sum "$MKVFILE" | cut -d" " -f1)
if [ $OLDFILESIZE -ne $NEWFILESIZE ]; then
error "ERROR: '$NEWFILE' and '$MKVFILE' files do not match. You might want to investigate!"
fi
fi
fi
# Remove new file in $WD
doprint "\nRemove working file."
doprint "> rm -f \"$NEWFILE\""
dopause
cleanup $NEWFILE
fi

timestamp "Copy file over took: "

# Run through the timestamp function manually to display total execution time
END=$(date +%s)
secs=$(($END - $START))
h=$(( secs / 3600 ))
m=$(( ( secs / 60 ) % 60 ))
s=$(( secs % 60 ))

if [ $EXECUTE = 1 -a $PAUSE = 0 ];then
color GREEN
echo -n "Total processing time: "
printf "%02d:%02d:%02d " $h $m $s
echo "($secs seconds)"
color OFF
echo
fi

NEWFILESIZE=$($DUCMD "$MKVFILE" | cut -f1) # NEWFILESIZE isn't available in some circumstances so just grab it again

# Print final filesize summary
if [ $EXECUTE = 1 -a $PAUSE = 0 ];then
color YELLOW; printf "Filesize summary:\n"; color OFF
printf "%23s %15d KB\n" "Original Filesize:" $MKVFILESIZE|sed -e :a -e 's/\(.*[0-9]\)\([0-9]\{3\}\)/\1,\2/;ta'
printf "%23s %15d KB\n" "Extracted DTS Filesize:" $DTSFILESIZE|sed -e :a -e 's/\(.*[0-9]\)\([0-9]\{3\}\)/\1,\2/;ta'
printf "%23s %15d KB\n" "Converted AC3 Filesize:" $AC3FILESIZE|sed -e :a -e 's/\(.*[0-9]\)\([0-9]\{3\}\)/\1,\2/;ta'
printf "%23s %15d KB\n" "Final Filesize:" $NEWFILESIZE|sed -e :a -e 's/\(.*[0-9]\)\([0-9]\{3\}\)/\1,\2/;ta'
fi