#!/bin/sh

## TODO: handle srt encoding?

_printhelp()
{
    cat <<EOF
Usage: ${1##*/} [OPTIONS] FILES|FOLDERS

Transcode FILES or files found in FOLDERS to .mkv with x264 and ogg. Black
stripes are *not* cropped by default since it may be sometimes inaccurate. By
default output filenames are the same as the original, with time appended.

Options:
  -b:      Default bitrate for audio stream with unidentifiable bitrate. If 0,
           copy stream (default).
  -c:      Copy streams (no reencoding).
  -C:      Enable auto-crop (needs video reencoding).
  -f:      Remove source when done.
  -h:      Display this help.
  -o OPT:  Additional options.
  -p:      Preview changes, do not encode.
  -P:      Generate two thumbnails, one cropped, the other not. Enable
           auto-crop, do not encode.
  -s:      Sample of 5 minutes.
  -S MIN:  Sample of MIN minutes.
  -t:      Remove all "title" metadata.

User options are read from variable TC_VIDEO_OPT. This can be useful if you
often use the same options.

Examples:

* Exchange stream 1 an 2 by first removing them, then reading them in the
  desired order.

  ${1##*/} -o '-map -0:1 -map -0:2 -map 0:2 -map 0:1' input.mkv

* Change audio stream 1 title and remove audio stream 2 title:

  ${1##*/} -o '-metadata:s:a:0 title="FR: OGG Stereo" -metadata:s:a:1 title=' input.mkv

EOF
}

SAMPLE=""
VIDEO_FILTER=""
VIDEO_PARAM="-c:v libx264 -preset slow -crf 20"
AUDIO_PARAM=""
AUDIO_CODEC="libvorbis"
AUDIO_DEFAULT_RATE=0

## What to do if file exists:
# -y overwrite
# -n do not overwrite.
OVERWRITE="-n"
OPT_OVERWRITE=false

OPT_REMOVE_TITLE=false
OPT_CROP=false
OPT_CROPPREVIEW=false
OPT_PREVIEW=false
OPT_COPY=false

while getopts ":b:cCfho:pPsS:t" opt; do
    case $opt in
        b)
            AUDIO_DEFAULT_RATE=$OPTARG ;;
        c)
            VIDEO_PARAM="-c:v copy"
            AUDIO_PARAM="-c:a copy"
            OPT_COPY=true;;
        C)
            OPT_CROP=true;;
        f)
            OVERWRITE="-y"
            OPT_OVERWRITE=true;;
        h)
            _printhelp "$0"
            exit 1 ;;
        o)
            TC_VIDEO_OPT="$OPTARG" ;;
        p)
            OPT_PREVIEW=true ;;
        P)
            OPT_CROP=true
            OPT_PREVIEW=true
            OPT_CROPPREVIEW=true ;;
        s)
            SAMPLE="-ss 60 -t 360" ;;
        S)
            SAMPLE="-ss 60 -t $((60*$OPTARG))" ;;
        t)
            OPT_REMOVE_TITLE=true ;;
        ?)
            _printhelp "$0"
            exit 1 ;;
        :)
            echo "Missing argument."
            _printhelp "$0"
            exit 1 ;;
    esac
done

shift $(($OPTIND - 1))
if [ $# -eq 0 ]; then
    _printhelp "$0"
    exit
fi

if ! command -v ffmpeg >/dev/null; then
    echo "ffmpeg required."
    exit
fi

_duration()
{
    ## In seconds.
    ffmpeg -nostdin -i "$1" 2>&1 | awk '/Duration/ {split($2, time, /:|\./); print time[1]*60*60 + time[2]*60 + time[3]}'
}

_highfreq()
{
    awk 'BEGIN{max=0} /crop=/ {t[$NF]++; if (t[$NF]>max) {max=t[$NF]; val=$NF}} END{print val}'
}

_cropvalue()
{
    ## For 5 different timeslices of 1 second at every 1/6th of the video, we
    ## sample the crop values. We keep the values with highest frequency.
    STEP=$(($(_duration "$1")/6))
    for i in $(seq $STEP $STEP $((5*$STEP))); do
        ffmpeg -nostdin -ss $i -i "$1"  -t 10 -vf "cropdetect=24:2:0" -f null - 2>&1
    done | _highfreq
}

_croppreview()
{
    STEP=$(($(_duration "$1")/6))
    for i in $(seq  $STEP $STEP $((5*$STEP))); do
            ffmpeg -nostdin -v warning -y -ss $i -i "$1" \
                -f image2 -vframes 1 $VIDEO_FILTER "${1%.*}-preview-$i-cropped.png" \
                -f image2 -vframes 1 "${1%.*}-preview-$i-uncropped.png"
    done
}

_audiobitrate()
{
    ## As of 2013-10-23, ffmpeg sucks at retrieving audio bitrate. Therefore we
    ## use mediainfo if available, ffmpeg otherwise.

    ## In the end, if some stream bitrates are missing, we have two
    ## possibilities: either we copy the stream (default), or we encode it to a
    ## default value.

    ## AWK note: in "(bitrate+0 > 500)", the +0 is required to make sure bitrate is numeric.

    ## WARNING: mediainfo puts a space every 3 digits, hence the 'for' loop.
    if command -v mediainfo >/dev/null; then
        mediainfo "$1" | awk 'BEGIN {id=0} /^Audio/ { getline; while(getline && ! index($0,"Bit rate  ")); for(i=4; i<NF; i++) res=$i; print id, res; id++}'
    else
        ffmpeg -nostdin -i "$1" 2>&1 | awk 'BEGIN {id=0} /^ *Stream.*Audio/ {match($0, / [^ ]+ kb\/s/); res=substr($0, RSTART+1, RLENGTH-6); print id, res; id++}'
    fi | awk -v codec="$AUDIO_CODEC" -v bitrate="$AUDIO_DEFAULT_RATE" \
        '{if($2 || bitrate) {if($2) bitrate=$2; if(bitrate+0 > 500) bitrate=500; printf "-c:a:" $1 " " codec " -b:a:" $1 " " bitrate "k "} else printf "-c:a:" $1 " copy "} END {print ""}'
}

_transcode()
{
    echo "==> [$1]"
    OUTPUT="${1%.*}-$(date '+%F-%H%M%S').mkv"

    STREAM_TITLE=""
    if $OPT_REMOVE_TITLE; then
        STREAM_NUM=$(ffmpeg -nostdin -i "$1" 2>&1 | grep -c 'Stream')
        for i in $(seq 0 $STREAM_NUM); do
            STREAM_TITLE="${STREAM_TITLE}-metadata:s:$i title= "
        done
    fi

    if $OPT_CROP; then
        echo "Computing crop values... "
        VIDEO_FILTER="-vf $(_cropvalue "$1")"
        if $OPT_CROPPREVIEW; then
            echo "Generating preview... "
            _croppreview "$1"
        fi
    fi

    ## WARNING: we mix down audio to 2 channels with '-ac 2'. This greatly
    ## reduce file size and avoid any confusion for playback, which is often the
    ## case when converting DTS to any other format because DTS has embedded
    ## channel description which is not available in these formats.
    ! $OPT_COPY && AUDIO_PARAM="$(_audiobitrate "$1") -ac 2"

    cat<<EOF
================================================================================
User options:  ${TC_VIDEO_OPT:-none}
Sample:        ${SAMPLE:-no}
Clear tags:    $OPT_REMOVE_TITLE
In place:      $OPT_OVERWRITE
Crop:          $OPT_CROP
Video param:   $VIDEO_PARAM
Video filter:  $VIDEO_FILTER
Audio param:   $AUDIO_PARAM
================================================================================
EOF

    $OPT_PREVIEW && return

    ## Zsh compatibility. We need it otherwise word splitting of parameter like
    ## SAMPLE will not work.
    STATUS="$(set -o | grep 'shwordsplit' | awk '{print $2}')"
    [ "$STATUS" = "off" ] && set -o shwordsplit

    ffmpeg -nostdin $OVERWRITE -i "$1" \
        $VIDEO_PARAM $VIDEO_FILTER \
        $AUDIO_PARAM \
        -c:s copy \
        -map 0 $STREAM_TITLE \
        $SAMPLE $TC_VIDEO_OPT "$OUTPUT"

    ## Restore Zsh previous options. This will not turn off shwordsplit if it
    ## was on before calling the function.
    [ "$STATUS" = "off" ] && set +o shwordsplit

    if $OPT_OVERWRITE; then
        rm "$1"
        mv -f "$OUTPUT" "${1%.*}.mkv"
    fi
    echo
}

for i in "$@"; do

    ## Argument is a folder. We search for all video files in there.
    if [ -d "$i" ]; then

        ## WARNING: ffmpeg continues to read stdin once it has started, so it should
        ## not be called from within a while<<EOF loop without disabling stdin.
        while IFS= read -r j; do
            _transcode "$j"
            done<<EOF
$(find "$i"  \( \
            -iname '*.avi' -o \
            -iname '*.flv' -o \
            -iname '*.mkv' -o \
            -iname '*.mov' -o \
            -iname '*.mp4' -o \
            -iname '*.mpeg' -o \
            -iname '*.mpg' -o \
            -iname '*.webm' -o \
            -iname '*.wmv' \) )
EOF

    else
        ## Argument is a regular file.
        _transcode "$i"
    fi

done