#!/bin/bash

# MYTHCUTKEY    Version 0.11
# Change log
# 0.11  cutlist: first cut piece was ignored (always included)
# seektable/duration wrong: simplify with mythcommflag & mythutil
# searches for near keyframe
#
# 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.
#
# ts stream audio & video can be out of sync by seconds. start of keep piece could be better 2 sec before cut point.

OVERWRITE=0
OUTPUT_DIR=""
CHANID=-9999
STARTTIME=""
rec_dir=""
basename=""
USER="mythtv"
PASS="mythtv" # Change this to "-p your_password" if necessary
USERPASSTABLE="-u$USER -p$PASS mythconverg"
UNDO=0
UNDO_DIR="/tmp/"
STARTSECS=`date +%s`
MARK_CUT_END=0
MARK_CUT_START=1

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 $USERPASSTABLE `
        basename=`echo "select basename from recorded where chanid=$CHANID and starttime=\"$STARTTIME\";" \
                | mysql -N $USERPASSTABLE `
#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 $USERPASSTABLE
        echo "load data infile \"$UNDO_DIR$basename.seek$num.sql\" into table recordedseek;" | mysql $USERPASSTABLE
        echo "delete from recordedmarkup where chanid=$CHANID and starttime=\"$STARTTIME\";" | mysql $USERPASSTABLE
        echo "load data infile \"$UNDO_DIR$basename.mark$num.sql\" into table recordedmarkup;" | mysql $USERPASSTABLE
        echo "delete from recorded where chanid=$CHANID and starttime=\"$STARTTIME\";" | mysql $USERPASSTABLE
        echo "load data infile \"$UNDO_DIR$basename.rec$num.sql\" into table recorded;" | mysql $USERPASSTABLE
        exit 0
fi

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

thefile=$rec_dir$basename
outfile=$OUTPUT_DIR$basename".new"
# MARKTYPES="0,1"
MARKTYPES="$MARK_CUT_END,$MARK_CUT_START"
# MARK_GOP_BYFRAME = 9
MARKKEYTYPES="6,7,9"

# Get the cuts from recordedmarkup table and store in an array
# returns zero indexed one dim array, 2 elements per seektable entry
# always returns [type, mark] (2 elements) as cuts[0] & cuts[1] if not blank..
cuts=( `echo "select type,mark from recordedmarkup where chanid=$CHANID and starttime=\"$STARTTIME\" \
                and type in ($MARKTYPES) order by mark" | mysql -N $USERPASSTABLE ` )

pieces=${#cuts[@]}
if [ $pieces == 0 ]; then
      echo "No cuts in seektable, aborting" >&2
      exit -1
fi

# Insert a cut mark record for the start of the recording
let k=pieces-1
while [ $k -gt 0 ]; do
        cuts[k+2]=${cuts[k]}
        cuts[k+1]=${cuts[k-1]}
        let k-=2
done
let pieces=pieces+2

# Cut'(1)= $MARK_CUT_END if first cut is not the start of recording
cuts[1]=1
if [ ${cuts[2]} -ne $MARK_CUT_END ]; then
        cuts[0]=$MARK_CUT_END
else
        cuts[0]=$MARK_CUT_START
fi
echo ${cuts[0]}  ${cuts[1]}  ${cuts[2]} ${cuts[3]}

# Check if last cut is not end of recording, i.e., we want to keep the end of the recording
if [ ${cuts[pieces-2]} == $MARK_CUT_END ]; then
        lastcut=${cuts[pieces-1]}
        lastseek=""
        lastseek=`echo "select mark from recordedseek where chanid=$CHANID and starttime=\"$STARTTIME\" and mark>$lastcut;" | \
                        mysql -N $USERPASSTABLE  | tail -n 1`
        if [ "$lastseek" != "" ]; then
                if [ $lastseek -gt $lastcut ]; then
                        cuts[pieces]=$MARK_CUT_START
                        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-1
i=0
k=0
nextmark=0
totaloffset=0
part=1
while [ $i -le $pieces ]; do
        if [ ${cuts[i]} -eq 0 ]; then # Look for a cut end, i.e., the start of a segment we want to keep
                # Find keyframes before or at keep piece
                cutstart=$((${cuts[i+1]}+1))
                cutlimit=$(($cutstart-200))

                key=( `echo "select mark,offset from recordedseek where mark>=$cutlimit and mark<=$cutstart and chanid=$CHANID and starttime=\"$STARTTIME\" \
                        and type in ($MARKKEYTYPES) order by offset;" | mysql -N $USERPASSTABLE | tail -n 1` )
                cutstart=$((${key[0]}))
                key1=$((${key[1]}))
                if [ $cutstart -le 1 ]; then
                        key1=0
                        cutstart=0
                fi

                # find keyframe after or at keep piece
                cutend=$((${cuts[i+3]}-2))
                lastseek=`echo "select mark from recordedseek where chanid=$CHANID and starttime=\"$STARTTIME\" order by mark desc limit 1;" | \
                        mysql -N $USERPASSTABLE `
                cutlimit=$(($lastseek))

                key=( `echo "select mark,offset from recordedseek where mark>=$cutend and mark<=$cutlimit and chanid=$CHANID and starttime=\"$STARTTIME\" \
                       and type in ($MARKKEYTYPES) order by -offset;" | mysql -N $USERPASSTABLE | tail -n 1` )

                echo "mark,offset = "${key[@]}" , lastseek= "$lastseek
                cutend=$((${key[0]}))
                key2=$((${key[1]}))

                # Check if cutend is past the last record in the seek table
                if [ $cutend -eq 0 ]; then
                        key2=`echo "select filesize from recorded where chanid=$CHANID and starttime=\"$STARTTIME\";" | \
                              mysql -N $USERPASSTABLE `
                        let cutend=$lastseek
                fi

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

                let k++

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

                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 $USERPASSTABLE
        echo "select * from recordedmarkup where chanid=$CHANID and starttime=\"$STARTTIME\" into outfile \"$UNDO_DIR$basename.mark$num.sql\";" | \
                mysql $USERPASSTABLE
        echo "select * from recorded where chanid=$CHANID and starttime=\"$STARTTIME\" into outfile \"$UNDO_DIR$basename.rec$num.sql\";" | \
                mysql $USERPASSTABLE

        ionice -c3 mythcommflag --rebuild --chanid $CHANID --starttime $STARTTIME
        mythutil --clearcutlist --chanid $CHANID --starttime $STARTTIME
        mythutil --clearskiplist --chanid $CHANID --starttime $STARTTIME

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"