#!/usr/bin/env bash

DESIRED_SUB_LANG=eng

# sets UNMAP_STREAMS
set_unmap_streams() {
    local file="$1"
    local unmapFilter='bin_data|jpeg|png'
    local streamsStr
    UNMAP_STREAMS=()
    streamsStr="$(get_num_streams "${file}")" || return 1
    mapfile -t streams <<<"${streamsStr}" || return 1
    for stream in "${streams[@]}"; do
        if [[ "$(get_stream_codec "${file}" "${stream}")" =~ ${unmapFilter} ]]; then
            UNMAP_STREAMS+=("-map" "-0:${stream}")
        fi
    done
}

# sets AUDIO_PARAMS
set_audio_params() {
    local file="$1"
    local videoLang
    AUDIO_PARAMS=()
    videoLang="$(get_stream_lang "${file}" 'v:0')" || return 1
    for stream in $(get_num_streams "${file}" 'a'); do
        local numChannels codec lang
        numChannels="$(get_num_audio_channels "${file}" "${stream}")" || return 1
        if [[ ${numChannels} == '' ]]; then
            echo_fail "could not obtain channel count for stream ${stream}"
            return 1
        fi
        local channelBitrate=$((numChannels * 64))
        codec="$(get_stream_codec "${file}" "${stream}")" || return 1
        lang="$(get_stream_lang "${file}" "${stream}")" || return 1
        if [[ ${videoLang} != '' && ${videoLang} != "${lang}" ]]; then
            AUDIO_PARAMS+=(
                '-map'
                "-0:${stream}"
            )
        elif [[ ${codec} == 'opus' ]]; then
            AUDIO_PARAMS+=(
                "-c:${OUTPUT_INDEX}"
                "copy"
            )
            OUTPUT_INDEX=$((OUTPUT_INDEX + 1))
        else
            AUDIO_PARAMS+=(
                "-filter:${OUTPUT_INDEX}"
                "aformat=channel_layouts=7.1|5.1|stereo|mono"
                "-c:${OUTPUT_INDEX}"
                "libopus"
                "-b:${OUTPUT_INDEX}"
                "${channelBitrate}k"
            )
            OUTPUT_INDEX=$((OUTPUT_INDEX + 1))
        fi
    done
}

# sets SUBTITLE_PARAMS
set_subtitle_params() {
    local file="$1"
    local convertCodec='eia_608'

    local defaultTextCodec
    if [[ ${SAME_CONTAINER} == false && ${OUTPUT} == *'.mkv' ]]; then
        defaultTextCodec='srt'
        convertCodec+='|mov_text'
    else
        defaultTextCodec='mov_text'
        convertCodec+='|srt'
    fi

    SUBTITLE_PARAMS=()
    for stream in $(get_num_streams "${file}" 's'); do
        local codec lang
        codec="$(get_stream_codec "${file}" "${stream}")" || return 1
        lang="$(get_stream_lang "${file}" "${stream}")" || return 1
        if [[ ${lang} != '' && ${DESIRED_SUB_LANG} != "${lang}" ]]; then
            SUBTITLE_PARAMS+=(
                '-map'
                "-0:${stream}"
            )
        elif [[ ${codec} =~ ${convertCodec} ]]; then
            SUBTITLE_PARAMS+=("-c:${OUTPUT_INDEX}" "${defaultTextCodec}")
            OUTPUT_INDEX=$((OUTPUT_INDEX + 1))
        elif [[ ${codec} == 'hdmv_pgs_subtitle' ]]; then
            PGS_SUB_STREAMS+=("${stream}")
            SUBTITLE_PARAMS+=(
                '-map'
                "-0:${stream}"
            )
        else
            # map -0 covers the stream but still want to increment the index
            OUTPUT_INDEX=$((OUTPUT_INDEX + 1))
        fi
    done
}

get_encode_versions() {
    action="${1:-}"

    ENCODE_VERSION="encode=$(git -C "${REPO_DIR}" rev-parse --short HEAD)"
    FFMPEG_VERSION=''
    VIDEO_ENC_VERSION=''
    AUDIO_ENC_VERSION=''

    # shellcheck disable=SC2155
    local output="$(ffmpeg -version 2>&1)"
    while read -r line; do
        if line_starts_with "${line}" 'ffmpeg='; then
            FFMPEG_VERSION="${line}"
        elif line_starts_with "${line}" 'libsvtav1'; then
            VIDEO_ENC_VERSION="${line}"
        elif line_starts_with "${line}" 'libopus='; then
            AUDIO_ENC_VERSION="${line}"
        fi
    done <<<"${output}"

    local version
    if [[ ${FFMPEG_VERSION} == '' ]]; then
        while read -r line; do
            if line_starts_with "${line}" 'ffmpeg version '; then
                read -r _ _ version _ <<<"${line}"
                FFMPEG_VERSION="ffmpeg=${version}"
                break
            fi
        done <<<"${output}"
    fi

    if [[ ${VIDEO_ENC_VERSION} == '' ]]; then
        version="$(get_pkgconfig_version SvtAv1Enc)"
        test "${version}" == '' && return 1
        VIDEO_ENC_VERSION="libsvtav1=${version}"
    fi

    if [[ ${AUDIO_ENC_VERSION} == '' ]]; then
        version="$(get_pkgconfig_version opus)"
        test "${version}" == '' && return 1
        AUDIO_ENC_VERSION="libopus=${version}"
    fi

    test "${FFMPEG_VERSION}" == '' && return 1
    test "${VIDEO_ENC_VERSION}" == '' && return 1
    test "${AUDIO_ENC_VERSION}" == '' && return 1

    if [[ ${action} == 'print' ]]; then
        echo "${ENCODE_VERSION}"
        echo "${FFMPEG_VERSION}"
        echo "${VIDEO_ENC_VERSION}"
        echo "${AUDIO_ENC_VERSION}"
    fi
    return 0
}

# given an input mkv/sup file,
# output a new mkv file with
# input metadata preserved
replace_mkv_sup() {
    local mkvIn="$1"
    local supIn="$2"
    local mkvOut="$3"
    local stream="${4:-0}"

    local json
    json="$(get_stream_json "${mkvIn}" "${stream}")" || return 1

    # x:y
    # x = stream json variable name
    # y = mkvmerge option name
    local optionMap=(
        disposition.default:--default-track-flag
        disposition.original:--original-flag
        disposition.comment:--commentary-flag
        disposition.forced:--forced-display-flag
        disposition.hearing_impaired:--hearing-impaired-flag
        tags.language:--language
        tags.title:--track-name
    )

    # start building mkvmerge command
    local mergeCmd=(
        mkvmerge
        -o "${mkvOut}"
    )
    for line in "${optionMap[@]}"; do
        IFS=: read -r key option <<<"${line}"
        local val
        val="$(jq -r ".streams[].${key}" <<<"${json}")" || return 1
        # skip undefined values
        test "${val}" == null && continue
        # always track 0 for single track
        mergeCmd+=("${option}" "0:${val}")
    done
    mergeCmd+=("${supIn}")

    "${mergeCmd[@]}"
}

crop_sup() {
    local inSup="$1"
    local outSup="$2"
    local left="$3"
    local top="$4"
    local right="$5"
    local bottom="$6"
    local warnMsg='Window is outside new screen area'
    local maxAcceptableWarn=5
    local offset=5

    # skip cropping if not needed
    if [[ "${left}${top}${right}${bottom}" == "0000" ]]; then
        cp "${inSup}" "${outSup}" || return 1
        return 0
    fi

    for ((try = 0; try < 30; try++)); do
        echo_info "cropping sup with ${left} ${top} ${right} ${bottom}"
        "${SUPMOVER}" \
            "${inSup}" \
            "${outSup}" \
            --crop \
            "${left}" "${top}" "${right}" "${bottom}" &>"${outSup}.out" || return 1
        # supmover does not error for out-of-bounds subtitles
        # so adjust crop value until there is most certainly no issue
        if [[ "$(grep -c "${warnMsg}" "${cropSup}.out")" -gt ${maxAcceptableWarn} ]]; then
            echo_warn "${warnMsg}, retrying... (try ${try})"
            test "${left}" -gt ${offset} && left=$((left - offset))
            test "${top}" -gt ${offset} && top=$((top - offset))
            test "${right}" -gt ${offset} && right=$((right - offset))
            test "${bottom}" -gt ${offset} && bottom=$((bottom - offset))
        else
            return 0
        fi
    done
    # if we got here, all tries were had, so indicate failure
    return 1
}

# extract PGS_SUB_STREAMS from INPUT
# and crop using CROP_VALUE
setup_pgs_mkv() {
    local pgsMkvOut="$1"

    if [[ ${#PGS_SUB_STREAMS[@]} -eq 0 ]]; then
        return
    fi

    check_for_supmover || return 1

    # setup tempdir
    local ogSup cropSup cropMkv tmpdir
    tmpdir="${pgsMkvOut}-dir"
    recreate_dir "${tmpdir}" || return 1

    # get video resolution
    local vidRes vidWidth vidHeight
    vidRes="$(get_resolution "${INPUT}")"
    IFS=x read -r vidWidth vidHeight <<<"${vidRes}"

    for stream in "${PGS_SUB_STREAMS[@]}"; do
        # extract sup from input
        ogSup="${tmpdir}/${stream}.sup"
        cropSup="${tmpdir}/${stream}-cropped.sup"
        cropMkv="${tmpdir}/${stream}.mkv"
        mkvextract "${INPUT}" tracks "${stream}:${ogSup}" || return 1

        # check sup resolution
        local supRes
        supRes="$(get_sup_resolution "${ogSup}")" || return 1
        local supWidth supHeight
        IFS=x read -r supWidth supHeight <<<"${supRes}"
        local left top right bottom
        # determine crop values
        # if the supfile is smaller than the video stream
        # crop using aspect ratio instead of resolution
        if [[ ${vidWidth} -gt ${supWidth} || ${vidHeight} -gt ${supHeight} ]]; then
            echo_warn "PGS sup (stream=${stream}) is somehow smaller than initial video stream"
            echo_warn "cropping based off of aspect ratio instead of resolution"
            left=0
            # (supHeight - ((vidHeight/vidWidth) * supWidth)) / 2
            top="$(awk '{ print int(($1 - ($2 / $3 * $4)) / 2) }' <<<"${supHeight} ${vidHeight} ${vidWidth} ${supWidth}")"
            right=${left}
            bottom=${top}
        # otherwise crop using the crop value
        elif [[ ${CROP_VALUE} != '' ]]; then
            # determine supmover crop based off of crop
            local res w h x y
            # extract ffmpeg crop value ("crop=w:h:x:y")
            IFS='=' read -r _ res <<<"${CROP_VALUE}"
            IFS=':' read -r w h x y <<<"${res}"

            # ffmpeg crop value
            # is different than supmover crop inputs
            left=${x}
            top=${y}
            right=$((supWidth - w - left))
            bottom=$((supHeight - h - top))
        # fallback to just the video resolution
        else
            left=$(((supWidth - vidWidth) / 2))
            top=$(((supHeight - vidHeight) / 2))
            right=${left}
            bottom=${top}
        fi

        if ! crop_sup "${ogSup}" "${cropSup}" "${left}" "${top}" "${right}" "${bottom}"; then
            rm -r "${tmpdir}" || return 1
            return 1
        fi

        if ! replace_mkv_sup "${INPUT}" "${cropSup}" "${cropMkv}" "${stream}"; then
            echo_fail "could not replace mkv sup for ${stream}"
            rm -r "${tmpdir}" || return 1
        fi
    done

    # merge all single mkv into one
    mkvmerge -o "${pgsMkvOut}" "${tmpdir}/"*.mkv
    local mergeRet=$?
    rm -r "${tmpdir}" || return 1
    return ${mergeRet}
}

encode_usage() {
    local OPT_INDS
    declare -A SHORT_OPTS LONG_OPTS DESC_OPTS
    parse_opt_map \
        OPT_INDS \
        SHORT_OPTS \
        LONG_OPTS \
        DESC_OPTS \
        "${ENCODE_OPT_MAP[@]}"

    echo "encode -i input [options] [output]"
    echo -e "\nOPTIONS:"
    for ind in "${OPT_INDS[@]}"; do
        echo -e "  ${SHORT_OPTS[${ind}]}, ${LONG_OPTS[${ind}]}${DESC_OPTS[${ind}]}"
    done
    echo -e "\n  [output] output filename (default: \${PWD}/av1-input-file-name.mkv)\n"
}

encode_update() {
    git -C "${REPO_DIR}" pull
}

set_encode_opts() {
    # default values
    PRESET=3
    CRF=25
    GRAIN=''
    CROP=false
    PRINT_OUT=false
    DV_TOGGLE=false
    ENCODE_INSTALL_PATH='/usr/local/bin/encode'
    SAME_CONTAINER="false"

    local ENCODE_OPT_MAP=(
        "-i --input input file"
        "-P --preset set preset (default ${PRESET})"
        "-C --crf set CRF (default ${CRF})"
        "-g --grain set film grain (default disabled)"
        "-p --print print the script instead of executing it"
        "-c --crop use crop detect to auto-crop"
        "-d --dv enable dolby vision (default ${DV_TOGGLE})"
        "-v --version print version info"
        "-s --same-container use same container as input, default is mkv"
        "-u --update update script (git pull ffmpeg-builder)"
        "-I --install system install at ${ENCODE_INSTALL_PATH}"
        "-U --uninstall uninstall from ${ENCODE_INSTALL_PATH}"
    )

    # only using -I/U
    local minOpt=1
    test $# -lt ${minOpt} && encode_usage && return 1

    local arg value
    while [[ $# -gt 0 ]]; do
        arg="$1"
        value="${2:-}"
        case "${arg}" in
        -i | --input)
            INPUT="${value}"
            shift 2
            ;;
        -P | --preset)
            if ! is_positive_integer "${value}"; then
                encode_usage
                return 1
            fi
            PRESET="${value}"
            shift 2
            ;;
        -C | --crf)
            if ! is_positive_integer "${value}" || test "${value}" -gt 63; then
                echo_fail "${value} is not a valid CRF value (0-63)"
                encode_usage
                return 1
            fi
            CRF="${value}"
            shift 2
            ;;
        -g | --grain)
            if ! is_positive_integer "${value}"; then
                encode_usage
                return 1
            fi
            GRAIN="film-grain=${value}:film-grain-denoise=1:adaptive-film-grain=1:"
            shift 2
            ;;
        -c | --crop)
            CROP=true
            shift
            ;;
        -p | --print)
            PRINT_OUT=true
            shift
            ;;
        -d | --dv)
            DV_TOGGLE=true
            shift
            ;;
        -v | --version)
            get_encode_versions print || return 1
            return "${FUNC_EXIT_SUCCESS}"
            ;;
        -s | --same-container)
            SAME_CONTAINER=true
            shift
            ;;
        -u | --update)
            encode_update || return 1
            return "${FUNC_EXIT_SUCCESS}"
            ;;
        -I | --install)
            echo_warn "attempting install"
            sudo ln -sf "${SCRIPT_DIR}/encode.sh" \
                "${ENCODE_INSTALL_PATH}" || return 1
            echo_pass "succesfull install"
            return "${FUNC_EXIT_SUCCESS}"
            ;;
        -U | --uninstall)
            echo_warn "attempting uninstall"
            sudo rm "${ENCODE_INSTALL_PATH}" || return 1
            echo_pass "succesfull uninstall"
            return "${FUNC_EXIT_SUCCESS}"
            ;;
        *)
            # OUTPUT will be the last (optional) arg
            if [[ $# -ne 1 ]]; then
                echo_fail "wrong flags given"
                encode_usage
                return 1
            fi
            OUTPUT="${arg}"
            shift
            ;;
        esac
    done

    # validate input
    if [[ -z ${INPUT} || ! -f ${INPUT} ]]; then
        echo_fail "input undefined or does not exist"
        encode_usage
        return 1
    fi

    # fallback output path
    if [[ -z ${OUTPUT} ]]; then
        local basename
        basename="$(bash_basename "${INPUT}")" || return 1
        OUTPUT="${PWD}/av1-${basename}"
    fi

    # use same container for output
    if [[ ${SAME_CONTAINER} == true ]]; then
        local fileFormat outputSuffix
        fileFormat="$(get_file_format "${INPUT}")" || return 1
        if [[ ${fileFormat} == 'MPEG-4' ]]; then
            outputSuffix='mp4'
        elif [[ ${fileFormat} == 'Matroska' ]]; then
            outputSuffix='mkv'
        else
            echo_fail "unrecognized input format"
            return 1
        fi
    else
        outputSuffix='mkv'
    fi
    OUTPUT="${OUTPUT%.*}"
    OUTPUT+=".${outputSuffix}"

    if [[ ${PRINT_OUT} == false ]]; then
        echo
        echo_info "INPUT: ${INPUT}"
        echo_info "GRAIN: ${GRAIN}"
        echo_info "OUTPUT: ${OUTPUT}"
        echo
    fi
}

# shellcheck disable=SC2034
# shellcheck disable=SC2155
# shellcheck disable=SC2016
gen_encode_script() {
    if missing_cmd mkvpropedit; then
        echo_fail "use: ${REPO_DIR}/scripts/install_deps.sh"
        return 1
    fi

    local outputBasename="$(bash_basename "${OUTPUT}")"
    local genScript="${TMP_DIR}/${outputBasename}.sh"

    # global output index number to increment
    OUTPUT_INDEX=0

    # single string params
    local params=(
        INPUT
        OUTPUT
        PRESET
        CRF
        CROP_VALUE
        ENCODE_VERSION
        FFMPEG_VERSION
        VIDEO_ENC_VERSION
        AUDIO_ENC_VERSION
        svtAv1Params
        pgsMkv
        muxxedPgsMkv
    )

    svtAv1ParamsArr=(
        "tune=0"
        "complex-hvs=1"
        "spy-rd=1"
        "psy-rd=1"
        "sharpness=3"
        "enable-overlays=0"
        "hbd-mds=1"
        "scd=1"
        "fast-decode=1"
        "enable-variance-boost=1"
        "enable-qm=1"
        "chroma-qm-min=10"
        "qm-min=4"
        "qm-max=15"
    )
    IFS=':'
    local svtAv1Params="${GRAIN}${svtAv1ParamsArr[*]}"
    unset IFS

    # arrays
    local arrays=(
        UNMAP_STREAMS
        AUDIO_PARAMS
        SUBTITLE_PARAMS
        videoParams
        metadata
        ffmpegParams
        PGS_SUB_STREAMS
    )
    local "${arrays[@]}"

    local videoParams=(
        "-crf" '${CRF}' "-preset" '${PRESET}'
    )
    local ffmpegParams=(
        '-hide_banner'
        '-i' '${INPUT}'
        '-y'
        '-map' '0'
        '-c:s' 'copy'
    )

    # set video params
    get_encode_versions || return 1
    local inputVideoCodec="$(get_stream_codec "${INPUT}" 'v:0')"
    if [[ ${inputVideoCodec} == 'av1' ]]; then
        ffmpegParams+=(
            "-c:v:${OUTPUT_INDEX}" 'copy'
        )
        # can't crop if copying codec
        CROP=false
    else
        ffmpegParams+=(
            '-pix_fmt' 'yuv420p10le'
            "-c:v:${OUTPUT_INDEX}" 'libsvtav1' '${videoParams[@]}'
            '-svtav1-params' '${svtAv1Params}'
        )
        metadata+=(
            '-metadata' '${VIDEO_ENC_VERSION}'
            '-metadata' 'svtav1_params=${svtAv1Params}'
            '-metadata' 'video_params=${videoParams[*]}'
        )
    fi
    OUTPUT_INDEX=$((OUTPUT_INDEX + 1))

    # these values may be empty
    set_unmap_streams "${INPUT}" || return 1
    set_audio_params "${INPUT}" || return 1
    set_subtitle_params "${INPUT}" || return 1

    if [[ ${UNMAP_STREAMS[*]} != '' ]]; then
        ffmpegParams+=('${UNMAP_STREAMS[@]}')
    fi

    if [[ ${AUDIO_PARAMS[*]} != '' ]]; then
        ffmpegParams+=('${AUDIO_PARAMS[@]}')
    fi

    if [[ ${SUBTITLE_PARAMS[*]} != '' ]]; then
        ffmpegParams+=('${SUBTITLE_PARAMS[@]}')
    fi

    metadata+=(
        '-metadata' '${ENCODE_VERSION}'
        '-metadata' '${FFMPEG_VERSION}'
    )

    # in the case all audio streams are copied,
    # don't add libopus metadata
    if line_contains "${AUDIO_PARAMS[*]}" 'libopus'; then
        metadata+=(
            '-metadata' '${AUDIO_ENC_VERSION}')
    fi

    local CROP_VALUE
    if [[ ${CROP} == true ]]; then
        CROP_VALUE="$(get_crop "${INPUT}")" || return 1
        ffmpegParams+=('-vf' '${CROP_VALUE}')
        metadata+=(
            '-metadata' '${CROP_VALUE}'
            '-metadata' "og_res=$(get_resolution "${INPUT}")"
        )
    fi

    # separate processing step for pkg subs
    local pgsMkv="${TMP_DIR}/pgs-${outputBasename// /.}.mkv"
    local muxxedPgsMkv='${OUTPUT}.muxxed'
    setup_pgs_mkv "${pgsMkv}" 1>&2 || return 1

    ffmpegParams+=('${metadata[@]}')

    {
        echo '#!/usr/bin/env bash'
        echo

        # add normal params
        for param in "${params[@]}"; do
            declare -n value="${param}"
            if [[ ${value} != '' ]]; then
                echo "${param}=\"${value[*]}\""
            fi
        done
        for arrName in "${arrays[@]}"; do
            declare -n arr="${arrName}"
            if [[ -v arr ]]; then
                echo "${arrName}=("
                printf '\t"%s"\n' "${arr[@]}"
                echo ')'
            fi
        done

        # actually do ffmpeg commmand
        echo
        if [[ ${DV_TOGGLE} == true ]]; then
            echo 'ffmpeg "${ffmpegParams[@]}" -dolbyvision 1 "${OUTPUT}" || \'
        fi
        echo 'ffmpeg "${ffmpegParams[@]}" -dolbyvision 0 "${OUTPUT}" || exit 1'

        # track-stats and clear title
        if [[ ${OUTPUT} == *'.mkv' ]]; then
            {
                # ffmpeg does not copy PGS subtitles without breaking them
                # use mkvmerge to extract and supmover to crop
                if [[ ${#PGS_SUB_STREAMS[@]} -gt 0 ]]; then
                    echo
                    echo 'mkvmerge -o "${muxxedPgsMkv}" "${pgsMkv}" "${OUTPUT}" || exit 1'
                    echo 'rm "${pgsMkv}" || exit 1'
                    echo 'mv "${muxxedPgsMkv}" "${OUTPUT}" || exit 1'
                fi

                echo 'mkvpropedit "${OUTPUT}" --add-track-statistics-tags || exit 1'
                echo 'mkvpropedit "${OUTPUT}" --edit info --set "title=" || exit 1'
            }
        fi

        echo
    } >"${genScript}"

    if [[ ${PRINT_OUT} == true ]]; then
        echo_info "${genScript} contents:" 1>&2
        echo "$(<"${genScript}")"
    else
        bash -x "${genScript}" || return 1
        rm "${genScript}"
    fi
}

FB_FUNC_NAMES+=('encode')
# shellcheck disable=SC2034
FB_FUNC_DESCS['encode']='encode a file using libsvtav1 and libopus'
encode() {
    # localize variables used by child functions
    local PRESET CRF GRAIN CROP PRINT_OUT \
        DV_TOGGLE ENCODE_INSTALL_PATH SAME_CONTAINER \
        INPUT OUTPUT

    set_encode_opts "$@"
    local ret=$?
    if [[ ${ret} -eq ${FUNC_EXIT_SUCCESS} ]]; then
        return 0
    elif [[ ${ret} -ne 0 ]]; then
        return ${ret}
    fi
    gen_encode_script || return 1
}