#!/usr/bin/env bash
# ---------------------------------------------------
#  youtube-dl GUI generating MKV file
#
#  Parameter :
#    * video stream page URL 
#  Depends on :
#    * zenity
#    * youtube-dl
#    * mkvtoolnix
#  Revision history :
#    20/02/2014, V1.0 - Creation by N. Bernaerts
#    10/03/2014, V1.1 - Progress with % and pulsate
#    22/03/2014, V1.2 - Add multi-stream download
#    01/04/2014, V1.3 - Handle empty attachments
#    05/03/2017, V1.4 - Code cleanup
# ---------------------------------------------------

# check youtube-dl and mkvmerge
command -v zenity >/dev/null 2>&1 || { echo "Please install zenity"; exit 1; }
command -v youtube-dl >/dev/null 2>&1 || { zenity --error --text="Please install youtube-dl"; exit 1; }
command -v mkvmerge >/dev/null 2>&1 || { zenity --error --text="Please install mkvmerge [mkvtoolnix package]"; exit 1; }

# check mkvmerge minimal version (v9.1.0)
MKV_VERSION=$(mkvmerge -V | cut -d' ' -f2 | tr -dc [0-9])
[ ${MKV_VERSION} -lt 910 ] && { zenity --error --text="mkvmerge should be v9.1.0 or higher"; exit 1; }

# config file
FILE_CONF="$HOME/.config/youtubedl-gui.conf"

# get default download directory
TARGET_DIR=$(grep "^directory=" "${FILE_CONF}" | cut -d'=' -f2-)

# if no target directory, default to home
[ "${TARGET_DIR}" = "" ] && TARGET_DIR="${HOME}"

# set separator as line feed
IFS=$'\n'

#  exit if script called without URL parameter
[ -z "$1" ] && STREAM_ERROR="Error : No URL passed"

# init variables and get stream main data
if [ -z "${STREAM_ERROR}" ]
then
  # create temporary directory & set log file
  TMP_DIR=$(mktemp -d)
  STREAM_LOG="${TMP_DIR}/youtubedl.log"

  # first parameter is stream URL
  STREAM_URL="$1"
  echo "## Video URL: ${STREAM_URL}" > "${STREAM_LOG}"

  {
  # get video title
  STREAM_TITLE=$(youtube-dl --get-title "${STREAM_URL}")
  [ "${STREAM_TITLE}" = "" ] && STREAM_TITLE="No title"
  echo "## Video Title: ${STREAM_TITLE}" >> "${STREAM_LOG}"

  # get available resolutions
  echo "## Getting video streams list" >> "${STREAM_LOG}"
  youtube-dl --list-formats "${STREAM_URL}" >> "${STREAM_LOG}"
  } | zenity --progress --pulsate --auto-close --width=700 --title "${STREAM_URL}" --text "Collecting stream title and available resolutions"
fi

# if no error, look for available streams 
if [ -z "${STREAM_ERROR}" ]
then
  # getting stream title previously logged
  STREAM_TITLE=$(grep "Video Title:" "${STREAM_LOG}" | sed 's/^.*: \(.*\)$/\1/')

  # extract stream table from the log file (all lines after the one starting with 'format ...
  STREAM_TABLE=$(sed -r -n -e '/^format/,${p}' "${STREAM_LOG}" | tail -n +2)

  # streams array generation : extract streams list from log, revert line order (best res. first) & extract one data per line
  STREAM_CHOICE=$(echo "${STREAM_TABLE}" | sed '1!G;h;$!d' | sed 's/^\([^ ]*\)[ ]*\([^ ]*\)[ ]*\([^ ]*\)[ ]*\(.*\)$/FALSE\n\1\n\3\n\2\n\4/')

  # if no stream detected by youtube-dl, error 
  [ "${STREAM_CHOICE}" == "" ] && STREAM_ERROR="No stream detected.\nThis site may not be compatible with youtube-dl." 
fi

# if no error, select stream(s) to download
if [ -z "${STREAM_ERROR}" ]
then
  # display stream selection dialog
  echo "## Selecting stream to download" >> "${STREAM_LOG}"
  TEXT="Please select stream to download.\nYou can select video only + audio only streams."
  STREAM_ID=$(zenity --list --checklist --title "${STREAM_TITLE}" --text "${TEXT}" --height=400 --width=700 --separator='+' --hide-column=2 --print-column=2 --column="Select" --column="Code" --column="Resolution" --column="Format" --column="Detail" ${STREAM_CHOICE})

  # handle zenity double answer bug #1267788 (in case)
  STREAM_ID="$(echo ${STREAM_ID} | sed 's/ /+/g' | cut -d'|' -f1)"
  echo "## Stream(s) selected : ${STREAM_ID}" >> "${STREAM_LOG}"

  # if no stream selected, error
  [ "${STREAM_ID}" = "" ] && STREAM_ERROR="Error : No stream selected"
fi

# if no error, start stream(s) download
if [ -z "${STREAM_ERROR}" ]
then
  # launch download command
  echo "## Stream(s) download - Beginning" >> "${STREAM_LOG}"
  youtube-dl --prefer-avconv -f "${STREAM_ID}" --write-thumbnail --write-description --all-subs --restrict-filenames --output "${TMP_DIR}/%(title)s-%(width)sx%(height)s.%(ext)s" "${STREAM_URL}" >> "${STREAM_LOG}" 2>&1 &

  # get command PID
  STREAM_PID=$(ps -ef | grep "youtube-dl" | grep "${STREAM_URL}" | sed 's/^'${USER}'[ ]*\([0-9]*\).*$/\1/g')
  echo "## Download process PID: ${STREAM_PID}" >> "${STREAM_LOG}"

  {
  # loop thru download progress (0 - 90%)
  STREAM_PERCENT=0
  STREAM_RUN=${STREAM_PID}
  while [ "${STREAM_RUN}" = "${STREAM_PID}" ]
  do
    # display download percentage, wait for 1 second and check if process is still running
    [ "${CURRENT_PROGRESS}" != "" ] && CURRENT_PROGRESS="(${CURRENT_PROGRESS} %)"
    echo "#Downloading stream ${STREAM_ID} ${CURRENT_PROGRESS}"
    echo "${STREAM_PERCENT}"
    sleep 1
    STREAM_RUN=$(ps aux | awk '{print $2 }' | grep "${STREAM_PID}")

    # get current download percentage
    CURRENT_PROGRESS=$(grep "\[download\]" "${STREAM_LOG}" | grep "%" | sed 's/^.* \([0-9.]*\)%.*$/\1/')
    CURRENT_PERCENT=$(echo "${CURRENT_PROGRESS}" | sed 's/\([0-9]*\)\..*/\1/')

    # if no percentage displayed in youtube-dl log, simulate pulsate by adding 2%, else convert % from 0-100 to 0-96
    [ "${CURRENT_PERCENT}" = "" ] && STREAM_PERCENT=$((${STREAM_PERCENT} + 2)) || STREAM_PERCENT=$((${CURRENT_PERCENT} * 96 / 100))

    # if pulse is up to 96%, back to 2%
    [ ${STREAM_PERCENT} -gt 96 ] && STREAM_PERCENT="2"
  done
  
  # end of stream download, check if error is looged
  echo "## Stream(s) download - End" >> "${STREAM_LOG}"
  LOG_ERROR=$(grep "ERROR" "${STREAM_LOG}")
  [ "${LOG_ERROR}" != "" ] && STREAM_ERROR="${LOG_ERROR}"
  
  # if no error, mux downloaded files to final MKV
  if [ -z "${STREAM_ERROR}" ]
  then
    # get parameters for final mkv muxing (96%)
    echo "96"
    echo "#Determining muxing parameters ..."
    echo "## Final video - Determining muxing parameters" >> "${STREAM_LOG}"

    # get description, thumbnail & subtitle filenames
    STREAM_DESCR=$(grep "Writing video description to:" "${STREAM_LOG}" | sed 's/^.*: \(.*\)$/\1/')
    echo "## Video description : ${STREAM_DESCR}" >> "${STREAM_LOG}"
    STREAM_THUMB=$(grep "Writing thumbnail to:" "${STREAM_LOG}" | sed 's/^.*: \(.*\)$/\1/')
    echo "## Video thumbnail : ${STREAM_THUMB}" >> "${STREAM_LOG}"
    STREAM_SUBS=$(grep "Writing video subtitles to:" "${STREAM_LOG}" | sed 's/^.*: \(.*\)$/\1/')
    echo "## Video subtitle : ${STREAM_SUBS}" >> "${STREAM_LOG}"

    # try to get stream filename from ffmpeg streams merging result
    STREAM_FILE=$(grep "Merging formats into" "${STREAM_LOG}" | sed 's/^.*into .\(.*\).$/\1/')

    # else, get filename from direct stream download
    [ "${STREAM_FILE}" == "" ] && STREAM_FILE=$(grep "Destination:" "${STREAM_LOG}" | sed 's/^.*: \(.*\)$/\1/')
    echo "## Video filename : ${STREAM_FILE}" >> "${STREAM_LOG}"

    # generate final MKV filename
    STREAM_MKV=$(echo ${STREAM_FILE%.*})".mkv"
   
    # MKV muxing parameters : filename & title
    MUX_ARR=("--output" "${STREAM_MKV}" "--title" "${STREAM_TITLE}")

    # MKV muxing parameters : thumbnail (if file exists and not empty)
    [ -s "${STREAM_THUMB}" ] && MUX_ARR=("${MUX_ARR[@]}" "--attachment-mime-type" "image/jpeg" "--attachment-name" "cover.jpg" "--attach-file" "${STREAM_THUMB}")

    # MKV muxing parameters : description (if file exists and not empty)
    [ -s "${STREAM_DESCR}" ] && MUX_ARR=("${MUX_ARR[@]}" "--attachment-mime-type" "text/plain" "--attachment-name" "comment.txt" "--attach-file" "${STREAM_DESCR}")

    # MKV muxing parameters : downloaded stream
    MUX_ARR=("${MUX_ARR[@]}" "${STREAM_FILE}")

    # MKV muxing parameters : subtitles (if file exists and not empty)
    [ -s "${STREAM_SUBS}" ] && MUX_ARR=("${MUX_ARR[@]}" "${STREAM_SUBS}")

    # muxing final MKV file (98%)
    echo "98"
    echo "#Muxing final MKV file ..."
    echo "## Final video : ${STREAM_MKV}" >> "${STREAM_LOG}"
    mkvmerge "${MUX_ARR[@]}" >> "${STREAM_LOG}"
    
    # move resulting file to target directory
    [ -s "${STREAM_MKV}" ] && mv "${STREAM_MKV}" "${TARGET_DIR}" || mv "${STREAM_FILE}" "${TARGET_DIR}"
  fi
  } | zenity --width=700 --progress --auto-close --title "${STREAM_TITLE}" 
fi

# if download has been canceled, kill main process and all its children
if [ -n "${STREAM_PID}" ]
then
  # get PID of running process
  STREAM_RUN=$(ps aux | awk '{print $2 }' | grep "${STREAM_PID}")

  # if process still running, kill it
  if [ "${STREAM_RUN}" == "${STREAM_PID}" ]
  then 
    pkill --parent ${STREAM_PID}
    pkill ${STREAM_PID}
    STREAM_ERROR="ERROR : Download has been canceled"
  fi
fi

# if no error, check if there is an error in final log
[ -z "${STREAM_ERROR}" ] && STREAM_ERROR=$(grep "ERROR" "${STREAM_LOG}")

# if no error, clean-up of all temporary files, else error message
[ "${STREAM_ERROR}" != "" ] && zenity --error --no-wrap --title "Download error" --text "${STREAM_ERROR}\n(temporary files in ${TMP_DIR} till you close this dialog)"

# remove temporary files
rm -R "${TMP_DIR}"