#!/bin/bash

log_file=
fail() {
  echo "$*" >&2
  if [[ -e ${log_file} ]]; then
    rm ${log_file}
  fi
  exit -1
}


usage() {
  cat << EOF
usage: normalize-video [OPTIONS] [FILES]

Options:

  -h        This help
  -s        Force software only encode and decode
  -o OLDER  Only process files older than OLDER. See 'find -newer'
  -f        Force transcode even if already meets normalized
            settings.

EOF
}

force=0
older_than=
sw_only=
while getopts fho:s opt; do
  case "${opt}" in
  s)
    sw_only=1
    ;;
  o)
    older_than="-and -not -newer ${OPTARG}"
    ;;
  f)
    force=1
    ;;
  h)
    usage
    exit 0
    ;;
  [?])
    usage
    exit -1
    ;;
  esac
done
shift $(( OPTIND-1 ))

VIDEO=$(getent group video | sed -E 's,^video:[^:]*:([^:]*):.*$,\1,')
RENDER=$(getent group render | sed -E 's,^render:[^:]*:([^:]*):.*$,\1,')

if [[ "${sw_only}" == "" ]]; then
  if [[ "${VIDEO}" != "" ]]; then
    sw_only=0
    ADD_GROUPS="--group-add ${VIDEO}"
    if [[ "${RENDER}" != "" ]]; then
      ADD_GROUPS+=" --group-add ${RENDER}"
    fi
  else
    sw_only=1
  fi
fi

quiet="
  -v quiet
  -loglevel error
"

# mstodate MILLISECONDS
#
# Given MILLISECONDS, convert to DAYS:HOURS:MINUTES:SECONDS.MS
mstodate() {
  scales=( "86400-d:" "3600-h:-02" "60-m:-02" "1-.-02" )
  echo $1 | sed -E 's/([[:digit:]]{3})$/ \1/' | while read sec msec; do
    for scale in ${scales[@]}; do
      parts=(${scale//-/ })
      divisor=${parts[0]}
      suffix=${parts[1]}
      min=${parts[2]}
      num_scale=$((sec / divisor))
      sec=$((sec - (num_scale * divisor)))
      printf "%${min}d%s" ${num_scale} ${suffix}
    done
    echo "${msec}s"
  done
}

# Works:
# docker run --device=/dev/dri --user=1000 --rm --group-add 44 --group-add 109 -v /multimedia/Downloads/Infinite.Storm.2022/test:/multimedia/Downloads/Infinite.Storm.2022/test -v /multimedia/Downloads/Infinite.Storm.2022/test:/multimedia/Downloads/Infinite.Storm.2022/test --entrypoint=ffmpeg intel-media-ffmpeg -v error -loglevel error -nostdin -hwaccel qsv -c:v h264_qsv -i /multimedia/Downloads/Infinite.Storm.2022/test/Infinite.Storm.2022.mkv -progress /dev/stdout -metadata title='Infinite Storm 2022' -vf 'scale_qsv=w=1920:h=802,hwdownload,format=nv12' -map 0 -map -0:d -c:a copy -c:s copy -c:v h264_qsv -preset veryslow -crf 20 -movflags +faststart /multimedia/Downloads/Infinite.Storm.2022/test/Infinite.Storm.2022.transcoded.mkv -y^

# Convert yuv420p10le to nv12
#docker run --device=/dev/dri --user=1000 --rm --group-add 44 --group-add 109 -v /multimedia/Downloads/Infinite.Storm.2022/test:/multimedia/Downloads/Infinite.Storm.2022/test -v /multimedia/Downloads/Infinite.Storm.2022/test:/multimedia/Downloads/Infinite.Storm.2022/test --entrypoint=ffmpeg intel-media-ffmpeg -v quiet -loglevel debug -nostdin -hwaccel qsv -resize 1920x1080 -i /multimedia/Downloads/Infinite.Storm.2022/test/Foundation.S01E02.Preparing.to.Live.mkv -progress /dev/stdout -metadata title='Foundation.S01E02.Preparing.to.Live' -pix_fmt nv12 -vf hwupload=extra_hw_frames=64,vpp_qsv=format=nv12 -map 0 -map -0:d -c:a copy -c:s copy -global_quality 20 -c:v h264_qsv -movflags +faststart /multimedia/Downloads/Infinite.Storm.2022/test/Foundation.S01E02.Preparing.to.Live.transcoded.mkv -y 2>&1
ffmpeg_hw_or_sw() {
  IN="${1}"
  shift
  OUT="${1}"
  shift
  TITLE="${1}"
  shift
  width=${1}
  shift
  height=${1}
  shift
  codec="${1}"
  shift
  pix_fmt="${1}"
  shift

  case "${pix_fmt}" in
  *p10le)
    sw_decode=1
    ;;
  *)
    sw_decode=0
    ;;
  esac

  DIN="$(dirname "${IN}")"
  DOUT="$(dirname "${OUT}")"

  if (( sw_only == 1 )); then
    echo "${PREFIX} Content will be transcoded to Matroska(yuv420p:h264:${width}x${height}) using software." >&2

    input_flags="
    "

    output_flags="
      -map 0
      -map -0:d
      -c:a copy
      -c:s copy
      -c:v libx264
      -max_muxing_queue_size 1024
      -vf scale=${width}:${height},format=yuv420p
      -preset veryslow
      -crf 20
    "
#      -metadata "title='${TITLE}'" \

    ffmpeg \
      ${quiet} \
      -nostdin \
      ${input_flags} \
      -i "${IN}" \
      -progress /dev/stdout \
      ${output_flags} \
      -movflags +faststart \
      "${OUT}" \
      -y
  else
    output_flags="
      -map 0
      -map -0:d
      -c:a copy
      -c:s copy
    "

    if (( sw_decode == 1 )); then
      echo "${PREFIX} Content will be transcoded from ${codec} to Matroska(yuv420p:h264:${width}x${height}) with sw decode:hw encode." >&2
      input_flags="
        -hwaccel qsv
      "
      output_flags="
        -pix_fmt nv12
        -vf scale=${width}:${height}
        -vf hwupload=extra_hw_frames=64,scale_qsv=w=${width}:h=${height},vpp_qsv=format=nv12
        ${output_flags}
        -global_quality 20
        -c:v h264_qsv
      "
    else
      if [[ "${codec}" =~ .*264.* ]] || [[ "${codec}" =~ avc* ]]; then
        echo "${PREFIX} Content will be transcoded from ${codec} to Matroska(yuv420p:h264:${width}x${height}) with hw decode:hw encode." >&2
        input_flags="
          -hwaccel qsv
          -c:v h264_qsv
        "
        output_flags="
          -vf scale_qsv=w=${width}:h=${height},hwdownload,format=nv12
          ${output_flags}
          -global_quality 20
          -c:v h264_qsv
        "
      elif [[ "${codec}" = .*265.* ]] || [[ "${codec}" =~ hev* ]]; then
        echo "${PREFIX} Content will be transcoded from ${codec} to Matroska(yuv420p:h264:${width}x${height}) with hw decode:hw encode."  >&2
        input_flags="
          -hwaccel qsv
          -c:v hevc_qsv
        "
        output_flags="
          -vf scale_qsv=w=${width}:h=${height},hwdownload,format=nv12
          ${output_flags}
          -global_quality 20
          -c:v h264_qsv
        "
      else
        echo "${PREFIX} Content will be transcoded from ${codec} to Matroska(yuv420p:h264:${width}x${height}) using software only." >&2
        input_flags="
        "
        output_flags="
          ${output_flags}
          -c:v libx264
          -vf scale=${width}:${height},format=yuv420p
          -preset veryslow
          -crf 20
        "
      fi
    fi

#        -metadata "title='${TITLE}'" \
    docker run \
      --device=/dev/dri \
      --user=$(id -u) \
      --rm \
      ${ADD_GROUPS} \
      -v "${DIN}":"${DIN}" \
      -v "${DOUT}":"${DOUT}" \
      --entrypoint=ffmpeg \
      intel-media-ffmpeg \
        ${quiet} \
        -nostdin \
        ${input_flags} \
        -i "${IN}" \
        -progress /dev/stdout \
        ${output_flags} \
        -movflags +faststart \
        "${OUT}" \
        -y
  fi
}

function move_to_backup {
  IN="$1"
  base="$(dirname "${IN}")"
  if [ ! -d "${base}" ]; then
    continue
  fi
  base=${BACKUP:-/multimedia/backup}"${base}"
  if [ ! -d "${base}" ]; then
    mkdir -p "${base}" || fail "Unable to mkdir '${base}'"
  fi
  if [ -d "${base}" ]; then
    echo -e "\nTransferring ${IN} to ${base}"
    rsync -avprlP --remove-source-files "${IN}" "${base}/" || fail "Unable to move '$IN'"
  fi
}

function check_and_convert {
  IN="$1"
  PREFIX="$2"

  width=$(ffprobe \
    -v error \
    -select_streams v:0 \
    -show_entries "stream=width" \
    "${IN}" | sed -ne 's,width=,,p'
  )

  height=$(ffprobe \
    -v error \
    -select_streams v:0 \
    -show_entries "stream=height" \
    "${IN}" | sed -ne 's,height=,,p'
  )

  pix_fmt=$(ffprobe \
    -v error \
    -select_streams v:0 \
    -show_entries "stream=pix_fmt" \
    "${IN}" | sed -ne 's,pix_fmt=,,p'
  )

  title=$(ffprobe -v error -show_entries format_tags=title -of default=noprint_wrappers=1 "${IN}" | sed -e s,^TAG:title=,,)
  if [[ "${title}" == "" ]]; then
    title=$(basename ${IN})
    title=${title%.???}
    title="${title//./ }"
  fi

  info=($(ffprobe -v error \
    -select_streams v:0 \
    -show_entries stream=codec_name:format=format_name \
    -of default=noprint_wrappers=1:nokey=1 "${IN}"))
  codec=${info[0]}
  container=${info[1]}

  update=0
  transcode=${force}
  if (( width > 1920 )); then
    echo "${PREFIX} Has a width of ${width} and will be scaled to 1920."
    height=$((((((1920*height/width))*2))/2))
    width=1920
    update=1
    transcode=1
  fi

  if [[ "${pix_fmt}" != "yuv420p" ]]; then
    echo "${PREFIX} Has pixel format of ${pix_fmt} and will be converted to yuv420p."
    update=1
    transcode=1
  fi

  if [[ "$codec" != "h264" ]]; then
    echo "${PREFIX} Is using ${codec} and will be converted to h264."
    update=1
    transcode=1
  fi

  if [[ ! "${container}" =~ mkv|matroska ]]; then
    echo "${PREFIX} Is contained in ${container} and will be repacked to matroska."
    update=1
  fi

  if (( update == 0 )) && (( force == 0 )); then
    echo "${PREFIX} Skipping ${IN} as it is already normalized."
    return 0;
  fi

  echo -n "${PREFIX} Counting frames in ${IN}: "
  FRAMES=$(ffprobe -v error -select_streams v:0 -show_entries stream -of default=noprint_wrappers=1 "${IN}" | sed -ne 's,^TAG:NUMBER_OF_FRAMES=,,p') || fail "Unable to count frames"
  if [[ "${FRAMES}" != "" ]]; then
    echo -e "\n${PREFIX} ${FRAMES} frames."
  else
    duration=$(ffprobe -v error -show_entries format=duration -of default=noprint_wrappers=1:nokey=1 "${IN}")
    fps=$(ffprobe -v error -select_streams v -of default=noprint_wrappers=1:nokey=1 -show_entries stream=r_frame_rate ${IN})
    FRAMES=$(printf "%.0f" $(bc <<< "${duration}*${fps}" ))
    echo -e "\n${PREFIX} ${FRAMES} (from duration ${duration}) frames."
  fi
  
  IN="$(realpath "${IN}")"
  OUT="${IN/%.???/.mkv}"
  ORIG="${OUT}"
  if [[ "${OUT}" == "${IN}" ]]; then
    OUT="${OUT/.mkv/.transcoded.mkv}"
  fi

  echo "${PREFIX} Output file is $(basename "${OUT}")."

  # Use stdbuf to flush stdout/stderr every line
  STARTTIME=$(date +%s)
  LASTTIME=${STARTTIME}
  LASTPOS=0

  log_file=$(mktemp)
  if (( transcode == 0 )); then
    echo "${PREFIX} Content will be repacked to Matroska(yuv420p:h264:${width}x${height})."

    input_flags="
    "

    output_flags="
      -map 0
      -map -0:d
      -c:a copy
      -c:s copy
      -c:v copy
    "
#-metadata title=\"${title}\" \

    command="
      ffmpeg \
        ${quiet} \
        -nostdin \
        ${input_flags} \
        -i \"${IN}\" \
        -progress /dev/stdout \
        ${output_flags} \
        -movflags +faststart \
        \"${OUT}\" \
        -y
      "
  else
    command="
      ffmpeg_hw_or_sw \
        \"${IN}\" \
        \"${OUT}\" \
        \"${title}\" \
        ${width} \
        ${height} \
        \"${codec}\" \
        \"${pix_fmt}\"
    "
  fi
  echo "Running: $command"
  { eval ${command} || {
      echo "${PREFIX} FFMPEG failed" | tee ${log_file}
      false
    }
  } | while read line; do
    if [[ "${line}" == "${PREFIX} FFMPEG failed" ]]; then
      cat ${log_file}
      fail "${PREFIX} FAIL: ${IN}"
      false
      break
    fi
    if [[ "${FRAMES}" == "" ]]; then
      frame=${frame:-$(echo $line | sed -n 's/^frame=*\(.*\)/Frame: \1/p')}
      speed=${speed:-$(echo $line | sed -n 's/^speeed=*\(.*\)/Speed: \1/p')}
      fps=${fps:-$(echo $line | sed -n 's/^fps=*\(.*\)/FPS: \1/p')}
      if [[ "${frame}" != "" ]] &&
        [[ "${speed}" != "" ]] &&
        [[ "${fps}" != "" ]]; then
        echo "${PREFIX} ${frame} ${speed} ${fps}"
        frame=
        fps=
        speed=
      else
        echo $line
      fi
      continue
    fi

    POS=$(echo $line | sed -n 's/^frame=*\(.*\)/\1/p')
    if [[ "${POS}" == "" ]]; then
      true
      continue
    fi

    NOW=$(date +%s)
    ELAPSEDTIME=$(( NOW - LASTTIME ))
    if (( ELAPSEDTIME <= 5 )); then
      true
      continue
    fi

    ELAPSEDFRAMES=$(( POS - LASTPOS ))
    REMAININGFRAMES=$(( FRAMES - POS ))
    REMAININGFRAMES=$(( REMAININGFRAMES * 1000 )) # convert to ms
    FRAMERATE=$(( ELAPSEDFRAMES / ELAPSEDTIME ))
    if (( FRAMERATE != 0 )); then
      REMAININGMS=$(( REMAININGFRAMES / FRAMERATE ))
      printf "\r%*s\r%s" $(tput cols) " " "${PREFIX} $(mstodate $(( NOW - STARTTIME ))000). Transcode $((ELAPSEDFRAMES / ELAPSEDTIME ))fps. Frame $POS of $FRAMES $(( $(( 100 * POS)) / FRAMES ))%. ETA $(mstodate ${REMAININGMS}) remaining."
    else
      printf "\r%*s\r%s" $(tput cols) " " "${PREFIX} $(mstodate $(( NOW - STARTTIME ))000). Transcode $((ELAPSEDFRAMES / ELAPSEDTIME ))fps. Frame $POS of $FRAMES $(( $(( 100 * POS)) / FRAMES ))%. No frames processed in last ${ELAPSEDTIME} seconds."
    fi
    LASTTIME=${NOW}
    LASTPOS=${POS}
  done || fail "${PREFIX} Unable to transcode"

  if grep -q "${PREFIX} FFMPEG failed" ${log_file}; then
    exit -1
  fi

  move_to_backup "${IN}"

  if [[ "${ORIG}" == "${IN}" ]]; then
    # If the original file was an '.mkv' then OUT was '.transcoded.mkv'
    # during transcode. Now that the transcode is complete, set the
    # file back to the original name.
    mv "${OUT}" "${IN}" || fail "${PREFIX} Unable to mv '${OUT}' -> '${IN}'"
  fi

  rm ${log_file}
  NOW=$(date +%s)
  echo -e "${PREFIX} Completed in $(mstodate $(( NOW - STARTTIME ))000)"
}

if [ -z "$1" ]; then
  find . -name '* *' -type d | while read file; do mv "${file}" "${file// /_}" || fail "Unable to rename dir ${file}"; done
  find . -name '* *' -type f | while read file; do mv "${file}" "${file// /_}" || fail "Unable to rename file ${file}"; done
  matches="
    \(
      -iname '*.mkv'
      -or -iname '*.mp4'
      -or -iname '*.avi'
      -or -iname '*.mov'
    \)
    -and -not -path './backup/*'
    ${older_than}
  "
  total=$(eval find . -type f ${matches} | wc -l)
  current=1
  eval find . -type f ${matches} | sort | while read file; do
    check_and_convert "${file}" "${current}/${total}:"
    current=$((current+1))
  done
else
  for file in "$*"; do
    check_and_convert "${file}" 1 1
  done
fi && {
  cat << EOF
To transfer the updated files:

cd /multimedia/Downloads
rsync -avprlP --remove-source-files /multimedia/Downloads/ azurite:/multimedia/
find . -type d -empty -delete

EOF
}