#!/bin/bash

# MYTHCUTKEY    Version 0.1
# This script cuts myth recordings at key frames using the MythTV seek table.
# Why? No external programs required, lossless and fast.
# Updates the myth database with sql calls including rebuilding the seek table.
# Output files may NOT be easily editable again and may have glitches at cut points.
#
# Undo option allows restoration of original recording and associated database
# records including seek table and markup table.
# NOTE: The undo function requires a directory for which mysql has write permissions.
# The default is /tmp/ however you will lose your undo information when you reboot.
# Undo file names are the recording file base name with seekN.sql, markN.sql or recN.sql
# appended, where N is a number starting from zero and the highest number indicating the
# latest undo files. The undo files are not huge but may eat up disk space over time
# after many edits.
#
# The script does not delete the original recording so you need disk free space equal
# to something less than the original recording file size depending on how much you cut
# from the file.

OVERWRITE=0
OUTPUT_DIR=""
CHANID=-9999
STARTTIME=""
rec_dir=""
basename=""
USER="root"
PASS="" # Change this to "-p your_password" if necessary
UNDO=0
UNDO_DIR="/tmp/"
STARTSECS=`date +%s`

USAGE="$0 -c channel_id -s starttime [-d dir] [-f filename] [-o directory] [-h] [-Z]
             [-u undo_dir] [-U]
        -h      Help/usage
        -c      Channel ID (required, %CHANID% from Mythtv)
        -s      Start time (required, %STARTTIME% from Mythtv)
        -d      Directory where the recording is store (optional, %DIR% from Mythtv)
        -f      Recording file name (optional, %FILE% from Mythtv)
        -o      Output directory (optional, recording directory if not specified)
        -Z      Overwrite original recording (optional, default is false, ignored if -o specified)
        -u      Directory where undo info is saved (optional, mysql must have permission to write
                to this directory, default is /tmp/)
        -U      Undo the last edit

        Example usage overwriting original recording (original is renamed, not deleted)
             mythcutkey -c %CHANID% -s %STARTTIME% -Z
        Example usage to save the file to a different directory
            mythcutnz -c %CHANID% -s %STARTTIME% -o /your/directory/
        Example usage specifying an undo directory
            mythcutnz -c %CHANID% -s %STARTTIME% -Z -u /myundodir
        Example usage to undo the last edit
            mythcutnz -c %CHANID% -s %STARTTIME% -U -u /myundodir

        Warning: The end files may not be editable and/or playable outside MythTV. Keep the original
        recording if you think you may want to do further edits or conversions in future. This script does
        NOT delete the original recording. It is renamed with the extension \".old\"."


while getopts "c:s:o:d:f:t:u:hZU" opt
do
        case $opt in
        c )
                CHANID="$OPTARG"
                ;;
        s )
                STARTTIME="$OPTARG"
                ;;
        d )
                rec_dir="$OPTARG"
                ;;
        f )
                basename="$OPTARG"
                ;;
        h )
                echo "Usage: $USAGE" >&2
                exit 0
                ;;
        o )
                OUTPUT_DIR="$OPTARG"
                ;;
        u )
                UNDO_DIR="$OPTARG"
                ;;
        Z )     OVERWRITE=1     # Will be reset to false if -o directory is specified
                ;;
        U )     UNDO=1
                ;;
        ? )
                echo "Invalid option: -$OPTARG" >&2
                echo "Usage: $USAGE" >&2
                exit -1
                ;;
        esac
done

if [ "$OUTPUT_DIR" != "" ]; then
        OVERWRITE=0
fi

if [ $CHANID == -9999 -o "$STARTTIME" == "" ]; then
        echo "Channel ID and/or Start Time missing" >&2
        echo "Usage: $USAGE" >&2
        exit -1
fi

#############################################################################
# Find where the recording is stored if not specified with -d and -f options
#############################################################################

if [ "$rec_dir" == "" -o "$basename" == "" ]; then
        storage_dirs=`echo "select dirname from storagegroup;" | mysql -N -u $USER $PASS mythconverg`
        basename=`echo "select basename from recorded where chanid=$CHANID and starttime=\"$STARTTIME\";" \
                | mysql -N -u $USER $PASS mythconverg`
#echo $basename
#echo $storage_dirs

        found=0
        for f in $storage_dirs; do
                if [ -e "$f$basename" ]; then
                        rec_dir=$f
                        found=1
                fi
        done
echo $rec_dir$basename

        if [ $found == 0 ]; then
                echo "Can't find recording, aborting" >&2
                exit -1
        fi
fi

# Set output directory to recording directory if not specified with -o option
if [ "$OUTPUT_DIR" == "" ]; then
        OUTPUT_DIR=$rec_dir
fi

if [ ! -d $UNDO_DIR ]; then
        echo "Can't find undo directory "$UNDO_DIR", aborting" >&2
        exit -1
fi

#############################################################################
# Undo the last edit if -U option specified
#############################################################################

if [ $UNDO == 1 ]; then
        if [ ! -e $rec_dir$basename".old" ]; then
                echo "Can't find original recording "$rec_dir$basename".old, aborting" >&2
                exit -1
        fi
        # Don't delete. Just rename.
        mv $rec_dir$basename $rec_dir$basename".new"
        mv $rec_dir$basename".old" $rec_dir$basename
        # Undo various database changes
        num=0
        while [ -e $UNDO_DIR$basename.seek$num.sql ]; do
                let num++
        done
        let num--
        echo "delete from recordedseek where chanid=$CHANID and starttime=\"$STARTTIME\";" | mysql -u $USER $PASS mythconverg
        echo "load data infile \"$UNDO_DIR$basename.seek$num.sql\" into table recordedseek;" | mysql -u $USER $PASS mythconverg
        echo "delete from recordedmarkup where chanid=$CHANID and starttime=\"$STARTTIME\";" | mysql -u $USER $PASS mythconverg
        echo "load data infile \"$UNDO_DIR$basename.mark$num.sql\" into table recordedmarkup;" | mysql -u $USER $PASS mythconverg
        echo "delete from recorded where chanid=$CHANID and starttime=\"$STARTTIME\";" | mysql -u $USER $PASS mythconverg
        echo "load data infile \"$UNDO_DIR$basename.rec$num.sql\" into table recorded;" | mysql -u $USER $PASS mythconverg
        exit 0
fi

#############################################################################
# Edit the recording based on cuts at keyframes
#############################################################################

thefile=$rec_dir$basename
outfile=$OUTPUT_DIR$basename".new"
MARKTYPES="0,1"

# Get the cuts from recordedmarkup table and store in an array
cuts=( `echo "select type,mark from recordedmarkup where chanid=$CHANID and starttime=\"$STARTTIME\" \
                and type in ($MARKTYPES) order by mark" | mysql -N -u $USER $PASS mythconverg` )

i=0
pieces=${#cuts[@]}

# Insert a record for the start of the recording if it does not exist
if [ ${cuts[1]} != 0 ]; then
        let k=pieces-1
        while [ $k -gt 0 ]; do
                cuts[k+2]=${cuts[k]}
                cuts[k+1]=${cuts[k-1]}
                let k-=2
        done
        cuts[0]=0
        cuts[1]=0
        let pieces=pieces+2
fi

# Check if last cut is not end of recording, i.e., we want to keep the end of the recording
if [ ${cuts[pieces-2]} == 0 ]; then
        lastcut=${cuts[pieces-1]}
        lastseek=""
        lastseek=`echo "select mark from recordedseek where chanid=$CHANID and starttime=\"$STARTTIME\" and mark>$lastcut;" | \
                        mysql -N -u $USER $PASS mythconverg | tail -n 1`
        if [ "$lastseek" != "" ]; then
                if [ $lastseek -gt $lastcut ]; then
                        cuts[pieces]=1
                        let cuts[pieces+1]=$lastseek+100        # Assume end of recording is within 100 frames of last seek table entry
                        let pieces=pieces+2
                fi
        fi
fi

let pieces=pieces-3
k=0
nextmark=0
totaloffset=0
part=1
while [ $i -le $pieces ]; do
        if [ ${cuts[i]} == 0 -o ${cuts[i]} == 5 ]; then # Look for a cut end, i.e., the start of a segment we want to keep
                # Cuts at keyframes are actually keyframe+1 so fix
                cutstart=$((${cuts[i+1]}-1))
                key1=""
                while [ "$key1" == "" ]; do
                        key1=`echo "select offset from recordedseek where mark=$cutstart and chanid=$CHANID and starttime=\"$STARTTIME\";" | \
                                mysql -N -u $USER $PASS mythconverg`
                        let cutstart--
                        if [ $cutstart -le 0 ]; then
                                key1=0
                                cutstart=-1
                        fi
                done
                cutend=$((${cuts[i+3]}-1))
                key2=""
                while [ "$key2" == "" ]; do
                        key2=`echo "select offset from recordedseek where mark=$cutend and chanid=$CHANID and starttime=\"$STARTTIME\";" | \
                                mysql -N -u $USER $PASS mythconverg`
                        # Check if cutend is past the last record in the seek table
                        lastseek=`echo "select mark from recordedseek where chanid=$CHANID and starttime=\"$STARTTIME\" order by mark desc limit 1;" | \
                                mysql -N -u $USER $PASS mythconverg`
                        if [ $cutend -gt $lastseek ]; then
                                key2=`echo "select filesize from recorded where chanid=$CHANID and starttime=\"$STARTTIME\";" | \
                                        mysql -N -u $USER $PASS mythconverg`
                                let key2--
                        else
                                let cutend++
                        fi
                done

                startmark[k]=$(($cutstart+1))
                endmark[k]=$(($cutend-1))
                markshift[k]=$(($cutstart+1-$nextmark))
                nextmark=$(($cutend-1-${markshift[k]})) # +1
                offsetshift[k]=$(($key1-$totaloffset))
                bytes=$(($key2-$key1))                  # +188
                totaloffset=$(($totaloffset+$bytes))

                let k++

                echo "Segment start="$key1" bytes (frame="$(($cutstart+1))", cutpoint frame was "$((${cuts[i+1]}-1))")"
                echo "Segment end="$key2" bytes (frame="$(($cutend-1))", cutpoint frame was "$((${cuts[i+3]}-1))")"

                kblocks=100
                blocksize=$((1024*$kblocks))
                blockstart=$((($key1+1)/$blocksize+1))
                pre=$(($blockstart*$blocksize-$key1))
                blocks=$((($bytes-$pre)/$blocksize))
                post=$(($bytes-$pre-$blocks*$blocksize))
                poststart=$((($blockstart+$blocks)*$blocksize))

                echo "Total bytes in segment="$bytes" ="$pre" bytes+"$blocks" blocks+"$post" bytes"

                # Cut and paste append the segment with dd
                echo "Cutting and pasting segment "$part"a"
                if [ $part == 1 ]; then
                        ionice -c3 dd ibs=1c obs=1K skip=$key1 count=$pre if=$thefile of=$outfile
                else
                        ionice -c3 dd ibs=1c obs=1K skip=$key1 count=$pre if=$thefile of=$outfile oflag=append conv=notrunc
                fi
                echo "Cutting and pasting segment "$part"b"
                ionice -c3 dd bs=$kblocks"K" skip=$blockstart count=$blocks if=$thefile of=$outfile oflag=append conv=notrunc
                if [ $post -gt 0 ]; then
                        echo "Cutting and pasting segment "$part"c"
                        ionice -c3 dd ibs=1c obs=1c skip=$poststart count=$post if=$thefile of=$outfile oflag=append conv=notrunc
                fi
                let part++
                # Increment loop counter by 4 since array has cut start/end pairs plus cut types
                let i+=4
        else
                # Skips the case where the first cut point is a cut start, i.e., when start of recording is deleted
                let i+=2
        fi
done
segs=$k

#############################################################################
# Reset the cut list, seek table, etc. for final file
#############################################################################

finalfile=$OUTPUT_DIR$basename
if [ $OVERWRITE == 1 ]; then
        # Don't delete original. Just rename.
        mv $rec_dir$basename $rec_dir$basename".old"
        mv $outfile $finalfile
        # Delete tmp file if it exists
        if [ -e $OUTPUT_DIR$basename".tmp" ]; then
                ionice -c3 rm $OUTPUT_DIR$basename".tmp"
        fi

        # Backup various data so we can undo the changes
        num=0
        while [ -e $UNDO_DIR$basename.seek$num.sql ]; do
                let num++
        done
        echo "select * from recordedseek where chanid=$CHANID and starttime=\"$STARTTIME\" into outfile \"$UNDO_DIR$basename.seek$num.sql\";" | \
                mysql -u $USER $PASS mythconverg
        echo "select * from recordedmarkup where chanid=$CHANID and starttime=\"$STARTTIME\" into outfile \"$UNDO_DIR$basename.mark$num.sql\";" | \
                mysql -u $USER $PASS mythconverg
        echo "select * from recorded where chanid=$CHANID and starttime=\"$STARTTIME\" into outfile \"$UNDO_DIR$basename.rec$num.sql\";" | \
                mysql -u $USER $PASS mythconverg

        # Clear the markup table
        echo "delete from recordedmarkup where chanid=$CHANID and starttime=\"$STARTTIME\";" | mysql -u $USER $PASS mythconverg
        # Clear bookmark, cutlist, commflagged flags
        echo "update recorded set bookmark=0, cutlist=0, commflagged=0 where chanid=$CHANID and starttime=\"$STARTTIME\";" | mysql -u $USER $PASS mythconverg
        # Update file size
        FILESIZE=`du -b $OUTPUT_DIR$basename | awk '{print $1}'`
        echo "update recorded set filesize=$FILESIZE where chanid=$CHANID and starttime=\"$STARTTIME\";" | mysql -u $USER $PASS mythconverg
        # Rebuild the seek table
        i=0
        while [ $i -lt $segs ]; do
                echo "Rebuilding seek table part "$i" of "$(($segs-1))
                if [ ${startmark[i]} -lt $((${endmark[i]}-${markshift[i]})) ]; then
                        delmark=${startmark[i]}
                else
                        delmark=$((${endmark[i]}-${markshift[i]}))
                fi
                echo "delete from recordedseek where chanid=$CHANID and starttime=\"$STARTTIME\" and \
                        mark>=$((${startmark[i]}-${markshift[i]})) and mark<$delmark;" | mysql -u $USER $PASS mythconverg
                echo "update recordedseek set mark=mark-${markshift[i]}, offset=offset-${offsetshift[i]} \
                        where chanid=$CHANID and starttime=\"$STARTTIME\" and mark>=${startmark[i]} and mark<${endmark[i]} order by mark;" | mysql -u $USER $PASS mythconverg
                let i++
        done
        echo "delete from recordedseek where mark>=$nextmark and chanid=$CHANID and starttime=\"$STARTTIME\";" | mysql -u $USER $PASS mythconverg
else
        # Don't accidentally overwrite an existing file
        if [ -e $finalfile ]; then
                i=1
                while [ -e $finalfile"."$i".mpg" ]; do
                        let i++
                done
                mv $outfile $finalfile"."$i".mpg"
        else
                mv $outfile $finalfile
        fi
fi

ENDSECS=`date +%s`
SECS=$(($ENDSECS-$STARTSECS))
MINS=$(($SECS/60))
SECS=$(($SECS-$MINS*60))
echo "Elapsed time: "$MINS" minutes "$SECS" seconds"