#!/bin/bash

function mkvSetTitle(){
    MKV="$1"
    KEY="$2"
    NFO="${KEY}.nfo"
    if [ -f "$NFO" ]; then
        TITLE=$(xmllint --xpath 'string(//title)' "$NFO")
        YEAR=$(xmllint --xpath 'string(//year)' "$NFO")
    fi

    if [ -n "$TITLE" ]; then 
        echo "Setting title to '$TITLE'"
        $RUN mkvpropedit --set title="$TITLE" "$MKV"
    fi
    # [ -n "$YEAR" ] && mkvpropedit --set year="$YEAR" "$MKV"
}

function mkvAddCover(){
    MKV="$1"
    KEY="$2"
    COVER="${KEY}-poster.jpg"

    if [ -f "$COVER" ]; then
        convert "$COVER" -resize 600x /tmp/cover.jpg
        echo "Attaching cover '$COVER'"
        $RUN mkvpropedit --replace-attachment name:cover.jpg:/tmp/cover.jpg "$MKV"
        if [ $? -ne 0 ]; then
            $RUN mkvpropedit --attachment-name cover.jpg --attachment-mime-type "image/jpeg" --add-attachment /tmp/cover.jpg "$MKV"
        fi
    fi
}

function mkvAddSmallCover(){
    MKV="$1"
    KEY="$2"
    COVER="${KEY}-poster.jpg"

    if [ -f "$COVER" ]; then
        convert "$COVER" -resize 120x /tmp/small_cover.jpg
        echo "Attaching small cover '$COVER'"
        $RUN mkvpropedit --replace-attachment name:small_cover.jpg:/tmp/small_cover.jpg "$MKV"
        if [ $? -ne 0 ]; then
            $RUN mkvpropedit --attachment-name small_cover.jpg --attachment-mime-type "image/jpeg" --add-attachment /tmp/small_cover.jpg "$MKV"
        fi
    fi
}

function mkvAddPoster(){
    MKV="$1"
    KEY="$2"
    FILE="${KEY}-poster.jpg"

    if [ -f "$FILE" ]; then
        echo "Attaching poster '$FILE'"
        $RUN mkvpropedit --replace-attachment "name:poster.jpg:$FILE" "$MKV"
        if [ $? -ne 0 ]; then
            $RUN mkvpropedit --attachment-name poster.jpg --attachment-mime-type "image/jpeg" --add-attachment "$FILE" "$MKV"
        fi
    fi
}

function mkvAddFanart(){
    MKV="$1"
    KEY="$2"
    FILE="${KEY}-fanart.jpg"

    if [ -f "$FILE" ]; then
        echo "Attaching fanart '$FILE'"
        $RUN mkvpropedit --replace-attachment "name:fanart.jpg:$FILE" "$MKV"
        if [ $? -ne 0 ]; then
            $RUN mkvpropedit --attachment-name fanart.jpg --attachment-mime-type "image/jpeg" --add-attachment "$FILE" "$MKV"
        fi
    fi
}

function mkvAddNfo(){
    MKV="$1"
    KEY="$2"
    FILE="${KEY}.nfo"

    if [ -f "$FILE" ]; then
        echo "Attaching NFO '$FILE'"
        $RUN mkvpropedit --replace-attachment "name:movie.nfo:$FILE" "$MKV"
        if [ $? -ne 0 ]; then
            $RUN mkvpropedit --attachment-name movie.nfo --attachment-mime-type "application/xml" --add-attachment "$FILE" "$MKV"
        fi
    fi
}

function mkvMake() {
    IN="$1"
    KEY="$2"
    EXT="${IN##*.}"
    if [ "$EXT" != "mkv" ]; then
        OUT="${KEY}.mkv"
        echo "Converting to mkv: $OUT".
        $RUN ffmpeg -i "$IN" -map 0 -c copy -use_wallclock_as_timestamps 1 -fflags +genpts -f matroska "$OUT"
    else
        echo "Already an MKV. $OUT"
    fi
}

function mkvAddAac() {
    MKV="$1"
    KEY="$2"
    ffprobe "$MKV" 2>&1 | grep Stream | grep dts
    if [ $? -eq 0 ]; then
        ffprobe "$MKV" 2>&1 | grep Stream | grep aac
        if [ $? -ne 0 ]; then
            OUT="${KEY}.AAC.mkv"
            ENCODE_AAC=""
            echo "DTS stream found, adding AAC: $OUT".
            $RUN ffmpeg -y -i "$MKV" -strict -2 -map 0 -c copy -map 0:a0 -c:a:1 aac -b:a 192k -ac 2 -f matroska "$OUT"
        fi
    fi
}

function mkvAddSrt() {
    MKV="$1"
    KEY="$2"
    SRT="${KEY}.srt"
    if [ -f "$SRT" ]; then
        OUT="${KEY}.SRT.mkv"
        echo "Adding SRT: $OUT".
        $RUN ffmpeg -y -i "$MKV" -i "$SRT" -map 0 -map 1 -c copy -metadata:s:s:1 language=eng  -f matroska "$OUT"
    fi
}


function mkvAddSub() {
    MKV="$1"
    KEY="$2"
    IDX="${KEY}.idx"
    SUB="${KEY}.sub"
    if [ -f "$SUB" ]; then
        OUT="${KEY}.SUB.mkv"
        echo "Adding SUB: $OUT".
        $RUN ffmpeg -y -i "$MKV" -i "$IDX" -i "$SUB"  -map 0 -c copy -map 1 -c:s:1 dvd_subtitle -f matroska "$OUT"
    fi
}

ARGS=$(getopt -o 'nas' -- "$@") || exit
eval "set -- $ARGS"

while true; do
    case $1 in
        -n)
            RUN=echo
            ;;
        -a)
            ADD_AAC=1
            ;;
        -s)
            ADD_SRT=1
            ;;
        --)
            shift
            break
            ;;
        *)
            exit 1
            ;;
    esac
    shift
done

INPUT="$1"
KEY="${INPUT%.*}"
OUT="$INPUT"
echo $INPUT
echo $KEY
echo $OUT

if [ -n "$ADD_AAC" ]; then
    mkvAddAac "$OUT" "$KEY"
    exit 0
fi

if [ -n "$ADD_SRT" ]; then
    mkvAddSrt "$OUT" "$KEY"
    mkvAddSub "$OUT" "$KEY"
    exit 0
fi

mkvMake "$OUT" "$KEY"
mkvSetTitle "$OUT" "$KEY"
mkvAddCover "$OUT" "$KEY"
mkvAddSmallCover "$OUT" "$KEY"
mkvAddPoster "$OUT" "$KEY"
mkvAddFanart "$OUT" "$KEY"
# mkvAddNfo "$OUT" "$KEY"