add-subtitles

#!/bin/bash

# SPDX-FileCopyrightText: 2017 - 2024 sudorook <daemon@nullcodon.com>
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
# This program is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
# for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program. If not, see <https://www.gnu.org/licenses/>.

set -Eeuo pipefail

ROOT="$(dirname "${0}")"

source "${ROOT}"/globals
source "${ROOT}"/langs

! check_command bc chardetect ffmpeg ffprobe find grep iconv sed uchardet && exit 3

#
# Globals
#

function detect_encoding {
  local res
  local encoding
  local confidence

  res="$(chardetect "${1}")"
  confidence="$(echo "${res}" | sed -n "s|.*with confidence \(.*\)$|\1|p")"

  # If chardetect confidence is too low, use file --mime-encoding to guess
  # character map instead.
  if [ -n "${confidence}" ] && [ "$(echo "${confidence} > .9" | bc -l)" -eq 1 ]; then
    encoding="$(echo "${res}" | sed -n "s|.*\ \([A-Za-z0-9\-]\+\) with .*|\1|p")"
  else
    show_warning "WARNING: Using fallback to get character encoding for ${1@Q}." >&2
    encoding="$(uchardet "${1}")"
    if [ -z "${encoding}" ]; then
      encoding="$(file -b --mime-encoding "${1}")"
    fi
  fi

  echo "${encoding}"
}

function convert_charset {
  local tmp
  local encoding

  tmp="$(mktemp)"
  encoding="$(detect_encoding "${1}")"

  # Use sed to delete the byte-order mark (BOM)
  if [[ "${encoding}" = UTF-8-SIG ]]; then
    # sed -e 's/\xef\xbb\xbf//' "${1}" > "${tmp}"
    sed -i 's/\xef\xbb\xbf//' "${1}"
    encoding="$(detect_encoding "${1}")"
  fi

  # Convert charset to UTF-8, if needed.
  case "${encoding}" in
    utf-8 | UTF-8 | ascii | us-ascii)
      # No need to do anything.
      ;;
    utf-8-sig | UTF-8-SIG)
      show_error "WARNING: Failed to remove byte-order mark in ${1@Q}. Skipping..."
      ;;
    unknown-8bit)
      show_warning "WARNING: Unknown character encoding for ${1@Q}. Skipping..."
      ;;
    binary) ;;
    *)
      show_info "INFO: Converting ${file@Q} from ${encoding@Q} to UTF-8." >&2
      iconv -f "${encoding}" -t utf-8 "${1}" > "${tmp}"
      mv "${tmp}" "${1}"
      ;;
  esac

  sync
}

function strip_webvtt_style {
  sed -i '/^STYLE
*$/{:a;N;N;/^\s*}
*$/M!ba;d}' "${1}"
}

function get_language {
  local in
  local name
  local end
  local iso_639_1_code
  local iso_639_1_name
  local iso_639_2_code
  local iso_639_2_name
  local iso_639_1_iso_3166_1_code
  local iso_639_1_iso_3166_1_code_strip
  local iso_639_1_iso_3166_1_name

  in="$(basename "${1}")"
  name="${in%.*}"

  # Single subtitles are usually English, so just return 'eng' if the video and
  # subtitle file names are the same.
  if [[ "${NAME}" = "${name}" ]] ||
     [[ "${NAME}" =~ ${name}( |-|_|\.)(sdh|SDH|hiHI)$ ]]; then
    show_warning "assuming ${in@Q} language is ${LANG_DEFAULT@Q}" >&2
    echo "${LANG_DEFAULT}"
    return
  fi

  # Get last string in the file name for checking whether it contains the
  # language of the subtitle stream.
  end="$(echo "${name,,}" |
         sed -E -e "s/[ -_]v?[[:digit:]]+$//g" \
                -e "s/[\ -_]sdh$|[\ -_]hi$//g" |
         grep -o '[^\ _]*$')"

  # Parse the end of the file name for three-character strings.
  for iso_639_2_code in "${!ISO_639_2[@]}"; do
    iso_639_2_name="${ISO_639_2[${iso_639_2_code}]}"
    if [[ "${end,,}" = "${iso_639_2_code,,}" ]] ||
       [[ "${end,,}" = "${iso_639_2_name,,}" ]]; then
      echo "${iso_639_2_code}"
      return
    fi
  done

  # Now, try parsing two-character strings.
  for iso_639_1_code in "${!ISO_639_1[@]}"; do
    if [[ "${end,,}" = "${iso_639_1_code,,}" ]]; then
      iso_639_1_name="${ISO_639_1[${iso_639_1_code}]}"

      # MP4 containers don't support ISO-639-1 codes, so map the 2-letter code
      # to the 3-letter ISO-639-2 one.
      if [[ "${EXTENSION}" = mp4 ]]; then
        for iso_639_2_code in "${!ISO_639_2[@]}"; do
          iso_639_2_name="${ISO_639_2[${iso_639_2_code}]}"
          if [[ "${iso_639_1_name}" = "${iso_639_2_name}" ]]; then
            echo "${iso_639_2_code}"
            return
          fi
        done
      else
        echo "${iso_639_1_code}"
        return
      fi
    fi
  done

  # Lastly, try parsing two-character strings + two-character country codes.
  for iso_639_1_iso_3166_1_code in "${!ISO_639_1_ISO_3166_1[@]}"; do
    iso_639_1_iso_3166_1_code_strip="${iso_639_1_iso_3166_1_code/[-_]/}"
    iso_639_1_iso_3166_1_name="${ISO_639_1_ISO_3166_1[${iso_639_1_iso_3166_1_code}]}"
    if [[ "${end,,}" = "${iso_639_1_iso_3166_1_code,,}" ]] ||
       [[ "${end,,}" = "${iso_639_1_iso_3166_1_code_strip,,}" ]] ||
       [[ "${end,,}" = "${iso_639_1_iso_3166_1_name,,}" ]]; then
      echo "${iso_639_1_iso_3166_1_code}"
      return
    fi
  done

  show_warning "${in@Q} defaulting to ${LANG_DEFAULT@Q}" >&2
  echo "${LANG_DEFAULT}"
}

function hardcode_cleanup {
  sed -i -e "s/\(\x92\|\x91\)/'/g" -e "s/\(\x94\|\x93\)/\"/g" "${1}"
}

function count_subs {
  local sub_count=0
  local stream_count
  local codec_type
  local idx
  local in="${1}"

  stream_count="$(ffprobe -v error -show_format file:"${in}" |
                  sed -n "s/nb_streams=\([0-9]\+\)/\1/p" 2> /dev/null)"
  stream_count="$((stream_count - 1))"
  for idx in $(seq 0 "${stream_count}"); do
    codec_type="$(ffprobe -v error -show_entries stream -select_streams "${idx}" file:"${in}" |
                  sed -n "s/^codec_type=\(.*\)/\1/p")"
    if [[ "${codec_type}" = subtitle ]]; then
      sub_count="$((sub_count + 1))"
    fi
  done

  echo "${sub_count}"
}

function find_subs {
  local subext
  local subenc
  local encoding
  local sub
  local subdir
  local subdirs=("Subs" "subs" "Subtitles" "subtitles")
  local file

  while read -r file; do
    if [ -f "${file}" ]; then
      subext=${file##*.}
      subenc="$(file -b --mime-encoding "${file}")"
      case "${subext}" in
        ssa | ass | srt | sup | smi)
          if [[ "${subenc}" != binary ]]; then
            convert_charset "${file}"
          fi
          SUBLIST+=("${file}")
          ;;
        vtt)
          if [[ "${subenc}" != binary ]]; then
            convert_charset "${file}"
            strip_webvtt_style "${file}"
          fi
          SUBLIST+=("${file}")
          ;;
        sub)
          if [ -f "${file%.*}.idx" ]; then
            IS_IMGSUB=true
          fi
          SUBLIST+=("${file}")
          ;;
        avi | m4v | mp4 | mkv | webm | idx) ;;
        *)
          show_warning "WARNING: ${file@Q} not supported. Skipping..."
          ;;
      esac
    fi
  done <<< "$(ls -v "${1}"*)"

  for subdir in "${subdirs[@]}"; do
    if [ -d "${subdir}" ]; then
      while read -r file; do
        if [ -f "${subdir}/${file}" ]; then
          subext=${file##*.}
          subenc="$(file -b --mime-encoding "${subdir}/${file}")"
          case "${subext}" in
            ssa | ass | srt | sup | smi)
              if [[ "${subenc}" != binary ]]; then
                convert_charset "${subdir}/${file}"
              fi
              SUBLIST+=("${subdir}/${file}")
              ;;
            vtt)
              if [[ "${subenc}" != binary ]]; then
                convert_charset "${file}"
                strip_webvtt_style "${file}"
              fi
              SUBLIST+=("${file}")
              ;;
            sub)
              if [ -f "${subdir}/${file%.*}.idx" ]; then
                IS_IMGSUB=true
              fi
              SUBLIST+=("${subdir}/${file}")
              ;;
            avi | m4v | mp4 | mkv | webm | idx) ;;
            *)
              show_warning "WARNING: ${file@Q} not supported. Skipping..."
              ;;
          esac
        fi
      done <<< "$(ls -v "${subdir}/")"
    fi
  done

  for sub in "${SUBLIST[@]}"; do
    show_listitem "${sub}"
  done
}

function get_codec {
  local fstring
  fstring="$(ffprobe -v error \
               -show_entries stream=codec_name -of compact=p=0:nk=1 file:"${1}")"
  echo "${fstring}"
}

function check_lang_match {
  local lang="${1}"
  local lang_len="${#1}"
  local lang_crop
  local default="${LANG_DEFAULT}"
  local default_len="${#default}"
  local default_crop

  if [ "${lang_len}" -eq "${default_len}" ]; then
    if [[ "${lang}" = "${default}" ]]; then
      return 0
    fi
  else
    case "${default_len}" in
      2)
        default_crop="${ISO_639_1[${default}]}"
        ;;
      3)
        default_crop="${ISO_639_2[${default}]}"
        ;;
      5)
        default_crop="$(echo "${ISO_639_1_ISO_3166_1[${default}]}" | cut -d" " -f1)"  # this could backfire...
        ;;
      *)
        show_error "ERROR: Something is wrong with default language ${default@Q}. Exiting..."
        exit 3
        ;;
    esac
    case "${lang_len}" in
      2)
        lang_crop="${ISO_639_1[${lang}]}"
        ;;
      3)
        lang_crop="${ISO_639_2[${lang}]}"
        ;;
      5)
        lang_crop="$(echo "${ISO_639_1_ISO_3166_1[${lang}]}" | cut -d" " -f1)"  # this could backfire...
        ;;
      *)
        show_error "ERROR: Something is wrong with subtitle language ${lang@Q}. Exiting..."
        exit 3
        ;;
    esac
    if [[ "${lang_crop}" = "${default_crop}" ]]; then
      return 0
    fi
  fi

  return 1
}

function make_cmd {
  local ext
  local sub
  local subidx
  local subname
  local subext
  local codec
  local inputs
  local mapping
  local ptsflag
  local metadata=""
  local copy
  local disposition
  local file_count
  local stream_count
  local lang
  local default_idx=-1 # idx of stream to set as the default

  ext="${CONTAINER:-"${EXTENSION}"}"

  inputs="-i file:${1@Q}"
  file_count=1
  mapping="-map_chapters 0 -map_metadata 0"
  if [[ "${FORCE}" = false ]]; then
    mapping="${mapping} -map 0"
    stream_count="$(count_subs "${1}")"
  else
    mapping="${mapping} -map 0:v -bsf:v 'filter_units=remove_types=6|39' -map 0:a"
    stream_count=0
  fi

  copy="-c:v copy -c:a copy"
  if [ "${stream_count}" -gt 0 ]; then
    copy="${copy}$(printf ' -c:s:%s copy' $(seq 0 $((stream_count - 1))))"
  fi

  for sub in "${SUBLIST[@]}"; do
    subname="${sub%.*}"
    subext="${sub##*.}"
    subidx="${sub%.*}.idx"
    lang="$(get_language "${sub}")"
    codec="$(get_codec "${sub}")"

    # add suffix if present (e.g. SDH for ambient sounds in the subtitles)
    if ! [ -f "${sub%.*}.idx" ]; then
      if [[ "${subname,,}" =~ ( |-|_|\.)sdh$ ]] ||
         [[ "${subname,,}" =~ ( |-|_|\.)hi$ ]]; then
        metadata="${metadata:+"${metadata}" }-metadata:s:s:${stream_count} language=${lang@Q}"
        if [ "${#lang}" = 3 ]; then
          metadata="${metadata} -metadata:s:s:${stream_count} handler=${ISO_639_2[${lang}]@Q}"
          metadata="${metadata} -metadata:s:s:${stream_count} title='${ISO_639_2[${lang}]} SDH'"
        elif [ "${#lang}" = 2 ]; then
          metadata="${metadata} -metadata:s:s:${stream_count} handler=${ISO_639_1[${lang}]@Q}"
          metadata="${metadata} -metadata:s:s:${stream_count} title='${ISO_639_1[${lang}]} SDH'"
        elif [ "${#lang}" = 5 ]; then
          metadata="${metadata} -metadata:s:s:${stream_count} handler=${ISO_639_1_ISO_3166_1[${lang}]@Q}"
          metadata="${metadata} -metadata:s:s:${stream_count} title='${ISO_639_1_ISO_3166_1[${lang}]} SDH'"
        fi
      else
        # only set default_idx if not sdh/hi to avoid ambient sounds in the
        # default subtitles
        if check_lang_match "${lang}"; then
          default_idx="${stream_count}"
        fi
        metadata="${metadata:+"${metadata}" }-metadata:s:s:${stream_count} language=${lang@Q}"
        if [ "${#lang}" = 3 ]; then
          metadata="${metadata} -metadata:s:s:${stream_count} handler=${ISO_639_2[${lang}]@Q}"
        elif [ "${#lang}" = 2 ]; then
          metadata="${metadata} -metadata:s:s:${stream_count} handler=${ISO_639_1[${lang}]@Q}"
        elif [ "${#lang}" = 5 ]; then
          metadata="${metadata} -metadata:s:s:${stream_count} handler=${ISO_639_1_ISO_3166_1[${lang}]@Q}"
        fi
      fi
    fi

    # set mappings and inputs based on the number of files being handed, and
    # then increment the stream and file counters
    if [[ "${subext}" = "sub" ]] && [ -f "${subidx}" ]; then
      local sub_stream_count
      inputs="${inputs} -i file:${subidx@Q} -i file:${sub@Q}"
      sub_stream_count="$(count_subs "${subidx}")"
      mapping="${mapping} -map ${file_count}:s"
      if [ "${sub_stream_count}" -gt 0 ]; then
        copy="${copy}$(printf ' -c:s:%s copy' $(seq "${stream_count}" $((stream_count + sub_stream_count - 1))))"
      fi

      # increment counters
      file_count=$((file_count + 2))
      stream_count=$((stream_count + sub_stream_count))
    else
      if [ "$(file -b --mime-encoding "${sub}")" = binary ]; then
        inputs="${inputs} -i file:${sub@Q}"
      else
        inputs="${inputs} -sub_charenc UTF-8 -i file:${sub@Q}"
      fi
      mapping="${mapping} -map ${file_count}:s"

      # set output subtitle codec
      if [[ "${codec}" = sami || "${codec}" = webvtt ]]; then
        case "${ext}" in
          mp4 | m4v)
            copy="${copy} -c:s:${stream_count} mov_text"
            ;;
          mkv)
            copy="${copy} -c:s:${stream_count} subrip"
            ;;
          *)
            copy="${copy} -c:s:${stream_count} copy"
            ;;
        esac
      elif [[ "${codec}" = subrip ]] && [[ "${ext}" = mp4 || "${ext}" = m4v ]]; then
        copy="${copy} -c:s:${stream_count} mov_text"
      elif [[ "${codec}" = microdvd ]] && [[ "${ext}" = mp4 || "${ext}" = m4v ]]; then
        copy="${copy} -c:s:${stream_count} mov_text"
      elif [[ "${codec}" =~ "96256" || "${codec}" =~ microdvd ]] &&
           [[ "${ext}" = mkv ]]; then
        copy="${copy} -c:s:${stream_count} srt"
      else
        copy="${copy} -c:s:${stream_count} copy"
      fi

      # increment counters
      file_count=$((file_count + 1))
      stream_count=$((stream_count + 1))
    fi
  done

  # copy any attachments
  copy="${copy} -map 0:t:? -c:t copy"

  if [[ "${NOSUB}" = false ]]; then
    if ! [[ "${default_idx}" == -1 ]]; then
      disposition="-disposition:s:${default_idx} default"
    else
      disposition=""
    fi
  else
    disposition="-default_mode infer_no_subs -disposition:s -default"
  fi

  if [[ "${CONTAINER}" != "${EXTENSION}" ]]; then
    ptsflag="-fflags +genpts"
  else
    ptsflag=""
  fi

  case ${ext} in
    avi)
      # AVI only supports embedding XSub files, so just use a Matroska
      # container instead.
      CONTAINER="mkv"
      show_warning "Re-muxing AVI to Matroska container." >&2
      ptsflag="-fflags +genpts"
      CMD="ffmpeg ${ptsflag} ${inputs} ${mapping} ${disposition} ${copy} -f matroska ${metadata} ${TMPVID@Q} -y"
      ;;
    mp4 | m4v)
      # MP4 can have issues embedding image-based subtitles, so convert to
      # Matroska.
      if "${IS_IMGSUB}"; then
        CONTAINER="mkv"
        ptsflag="-fflags +genpts"
        show_warning "Image subtitles detected. Re-muxing to Matroska container." >&2
        CMD="ffmpeg ${ptsflag} ${inputs} ${mapping} ${disposition} ${copy} -f matroska ${metadata} ${TMPVID@Q} -y"
      else
        CMD="ffmpeg ${ptsflag} ${inputs} ${mapping} ${disposition} ${copy} -f mp4 ${metadata} ${TMPVID@Q} -y"
      fi
      ;;
    mkv)
      CMD="ffmpeg ${ptsflag} ${inputs} ${mapping} ${disposition} ${copy} -f matroska ${metadata} ${TMPVID@Q} -y"
      ;;
    *)
      show_error "ERROR: ${ext@Q} files are not supported."
      exit 3
      ;;
  esac
}

function backup_files {
  local backupdir
  local infile
  local indir
  local subfile
  local subdir
  local subext
  local sub

  infile="$(basename "${IN}")"
  indir="$(dirname "${IN}")"
  backupdir="${indir}/backup_$(date +%Y%m%d-%H%M%S)"
  mkdir -p "${backupdir}"

  mv "${IN}" "${backupdir}/${infile}"
  for sub in "${SUBLIST[@]}"; do
    subdir="$(dirname "${sub}")"
    subfile="$(basename "${sub}")"
    subext=${sub##*.}
    if ! [[ "${subdir}" = "${indir}" ]]; then
      mkdir -p "${backupdir}/${subdir}"
    else
      subdir="."
    fi
    mv "${sub}" "${backupdir}/${subdir}/${subfile}"
    if [[ "${subext}" = sub ]]; then
      if [ -f "${sub%.*}.idx" ]; then
        mv "${sub%.*}.idx" "${backupdir}/${subdir}/${subfile%.*}.idx"
      fi
    fi
  done

  if [ -d "${indir}/Subs" ]; then
    rmdir "${indir}/Subs" --ignore-fail-on-non-empty
  fi
}

#
# Main
#

OPTIONS=c:fl:n
LONGOPTIONS=container:,force,language:,nodefault
PARSED=$(getopt -o "${OPTIONS}" --long "${LONGOPTIONS}" -n "${0}" -- "${@}")
eval set -- "${PARSED}"

FORCE=false
NOSUB=false
CONTAINER=
CMD=
IS_IMGSUB=false
LANG_DEFAULT=eng # needs to be 3-letter code (lower case)
while [ ${#} -ge 1 ]; do
  case ${1} in
    -c | --container)
      CONTAINER="${2}"
      shift 2
      ;;
    -f | --force)
      FORCE=true
      shift 1
      ;;
    -n | --nodefault)
      NOSUB=true
      shift 1
      ;;
    -l | --language)
      LANG_DEFAULT="${2}"
      shift 2
      ;;
    --)
      shift
      break
      ;;
    *)
      show_error "ERROR: invalid flag."
      exit 3
      ;;
  esac
done

if [[ ${#} -ne 1 ]]; then
  show_error "ERROR: please specify 1 input file."
  exit 3
fi

if ! [[ -f "${1}" ]]; then
  show_error "ERROR: input ${1@Q} not a file."
  exit 3
fi

IN="${1}"
NAME="${1%.*}"
EXTENSION="${1##*.}"
TMPVID="$(mktemp)"
SUBLIST=()

trap 'rm "${TMPVID}"; exit' INT TERM ERR

find_subs "${NAME}"
make_cmd "${IN}"
show_header "${CMD}"
eval "${CMD}"
backup_files
mv "${TMPVID}" "${NAME}.${CONTAINER:-"${EXTENSION}"}"
sync