#!/usr/bin/env bash
# ============================================================
#  Universal MKV → HEVC 8-bit MP4 Batch Encoder
#  Supports: AMD AMF, AMD VA-API, NVIDIA NVENC, Intel VA-API, Software (x265)
#  Audio:     First Japanese track only → AAC 128kbps stereo
#  Subtitles: Burned in with fallback priority:
#               1. English-tagged (eng)
#               2. Default-flagged subtitle in MKV
#               3. First available subtitle track
#               4. Skip file if no subtitles exist
#  Aspect:    4:3 source → 960x720 / 16:9 source → 1280x720
#  Chapters:  Preserved
#  Metadata:  Title set from filename
#
#  Requirements: ffmpeg + ffprobe must be in PATH
# ============================================================

set -euo pipefail

# --- SETTINGS -----------------------------------------------

QP=31
AUDIO_BITRATE=128k
OUTPUT_DIR="encoded_output"
SHARPEN="3:3:0.3"

# ============================================================

echo "============================================================"
echo " Universal HEVC 8-bit Encoder"
echo " QP=31  Audio=128kbps  4:3=960x720  16:9=1280x720"
echo "============================================================"
echo
echo " Select your GPU / encoder:"
echo
echo "   1. AMD      (hevc_amf   - AMF hardware, fastest)"
echo "   2. AMD      (hevc_vaapi - VA-API hardware, no extra install)"
echo "   3. NVIDIA   (hevc_nvenc - NVENC hardware)"
echo "   4. Intel    (hevc_vaapi - VA-API hardware, no extra install)"
echo "   5. Software (libx265    - slow but universal)"
echo
read -rp "  Enter choice (1-5): " GPU_CHOICE

# VA-API render device — /dev/dri/renderD128 covers most single-GPU systems.
# If you have both AMD and Intel GPUs, Intel is usually renderD129.
VAAPI_DEVICE="/dev/dri/renderD128"

case "$GPU_CHOICE" in
    1)
        ENCODER="hevc_amf"
        ENC_LABEL="AMD AMF"
        ENC_PARAMS="-rc cqp -qp_i $QP -qp_p $QP -quality quality -profile:v main -pix_fmt yuv420p"
        USE_VAAPI=0

        # Check for AMF runtime and offer to install if missing
        if ! ldconfig -p 2>/dev/null | grep -q "libamfrt64.so" && \
           ! find /usr/lib /usr/local/lib 2>/dev/null | grep -q "libamfrt64.so"; then
            echo
            echo " WARNING: AMF runtime (libamfrt64.so) not found."
            echo " The package 'amf-amdgpu-pro' is required from the AUR."
            echo
            read -rp " Install it now via yay? (y/n): " INSTALL_AMF
            if [[ "${INSTALL_AMF,,}" == "y" ]]; then
                echo
                echo " Installing amf-amdgpu-pro..."
                if ! command -v yay &>/dev/null; then
                    echo " ERROR: yay not found. Please install amf-amdgpu-pro manually."
                    exit 1
                fi
                yay -S --noconfirm amf-amdgpu-pro
                echo " AMF runtime installed."
            else
                echo
                echo " Aborting. Re-run the script and choose option 2 (VA-API) instead,"
                echo " which works without any additional packages."
                exit 1
            fi
        fi
        ;;
    2)
        ENCODER="hevc_vaapi"
        ENC_LABEL="AMD VA-API"
        ENC_PARAMS="-rc_mode CQP -qp $QP -profile:v main"
        USE_VAAPI=1
        ;;
    3)
        ENCODER="hevc_nvenc"
        ENC_LABEL="NVIDIA NVENC"
        ENC_PARAMS="-rc constqp -qp $QP -preset p4 -profile:v main -pix_fmt yuv420p"
        USE_VAAPI=0
        ;;
    4)
        ENCODER="hevc_vaapi"
        ENC_LABEL="Intel VA-API"
        ENC_PARAMS="-rc_mode CQP -qp $QP -profile:v main"
        USE_VAAPI=1
        ;;
    5)
        ENCODER="libx265"
        ENC_LABEL="Software x265"
        ENC_PARAMS="-crf 22 -preset slow -profile:v main -pix_fmt yuv420p"
        USE_VAAPI=0
        ;;
    *)
        echo
        echo " Invalid choice. Please enter 1, 2, 3, 4 or 5."
        exit 1
        ;;
esac

echo
echo " Using encoder: $ENC_LABEL ($ENCODER)"
echo " Output folder: $OUTPUT_DIR"
echo "============================================================"
echo

mkdir -p "$OUTPUT_DIR"

TMPFILE="$(mktemp /tmp/ffprobe_tmp.XXXXXX)"
trap 'rm -f "$TMPFILE"' EXIT

ENCODED=0
FAILED=0
SKIPPED=0

shopt -s nullglob
MKV_FILES=(*.mkv)

if [[ ${#MKV_FILES[@]} -eq 0 ]]; then
    echo " No .mkv files found in current directory."
    exit 0
fi

for FILEPATH in "${MKV_FILES[@]}"; do
    FILENAME="$FILEPATH"
    BASENAME="${FILEPATH%.mkv}"

    echo "-------------------------------------------------------"
    echo "[$FILENAME] Analysing streams..."

    # ================================================================
    # SUBTITLE SELECTION — three-tier fallback
    # ================================================================
    SUB_INDEX=""
    SUB_METHOD=""

    ffprobe -v error -select_streams s \
        -show_entries stream=index,disposition=default:stream_tags=language \
        -of csv=p=0 "$FILENAME" > "$TMPFILE" 2>&1

    # ---- Tier 1: English-tagged subtitle ----
    SUB_REL_COUNTER=0
    while IFS= read -r LINE; do
        if [[ -z "$SUB_INDEX" ]]; then
            if echo "$LINE" | grep -qi ",eng"; then
                SUB_INDEX="$SUB_REL_COUNTER"
                SUB_METHOD="English-tagged"
            fi
        fi
        (( SUB_REL_COUNTER++ )) || true
    done < "$TMPFILE"

    # ---- Tier 2: Default-flagged subtitle ----
    if [[ -z "$SUB_INDEX" ]]; then
        SUB_REL_COUNTER=0
        while IFS= read -r LINE; do
            if [[ -z "$SUB_INDEX" ]]; then
                if echo "$LINE" | grep -qP ",1,|,1$"; then
                    SUB_INDEX="$SUB_REL_COUNTER"
                    SUB_METHOD="default-flagged"
                fi
            fi
            (( SUB_REL_COUNTER++ )) || true
        done < "$TMPFILE"
    fi

    # ---- Tier 3: First available subtitle ----
    if [[ -z "$SUB_INDEX" ]] && [[ -s "$TMPFILE" ]]; then
        SUB_INDEX=0
        SUB_METHOD="first available"
    fi

    # ================================================================
    # Skip if no subtitles found
    # ================================================================
    if [[ -z "$SUB_INDEX" ]]; then
        echo "[$FILENAME] WARNING: No subtitle tracks found. Skipping file."
        (( SKIPPED++ )) || true
        continue
    fi

    echo "[$FILENAME] Subtitle selected: index $SUB_INDEX ($SUB_METHOD)"

    # ================================================================
    # RESOLUTION — 4:3 → 960x720, 16:9 → 1280x720
    # ================================================================
    STORED_W=""
    STORED_H=""
    SAR_N=1
    SAR_D=1

    STORED_W=$(ffprobe -v error -select_streams v:0 \
        -show_entries stream=width -of default=nw=1:nk=1 "$FILENAME" 2>/dev/null | head -1)

    STORED_H=$(ffprobe -v error -select_streams v:0 \
        -show_entries stream=height -of default=nw=1:nk=1 "$FILENAME" 2>/dev/null | head -1)

    SAR_RAW=$(ffprobe -v error -select_streams v:0 \
        -show_entries stream=sample_aspect_ratio -of default=nw=1:nk=1 "$FILENAME" 2>/dev/null | head -1)

    if [[ "$SAR_RAW" != "N/A" && "$SAR_RAW" =~ ^([0-9]+):([0-9]+)$ ]]; then
        SAR_N="${BASH_REMATCH[1]}"
        SAR_D="${BASH_REMATCH[2]}"
        [[ "$SAR_N" == "0" ]] && SAR_N=1
        [[ "$SAR_D" == "0" ]] && SAR_D=1
    fi

    DISP_W=$(( STORED_W * SAR_N / SAR_D ))

    SCALE_FILTER="scale=1280:720:flags=lanczos"
    RES_LABEL="1280x720 (16:9 default)"

    if [[ -n "$DISP_W" && -n "$STORED_H" && "$STORED_H" -gt 0 ]]; then
        RATIO10=$(( DISP_W * 10 / STORED_H ))
        if [[ "$RATIO10" -ge 14 ]]; then
            SCALE_FILTER="scale=1280:720:flags=lanczos"
            RES_LABEL="1280x720 (16:9)"
        else
            SCALE_FILTER="scale=960:720:flags=lanczos"
            RES_LABEL="960x720 (4:3)"
        fi
    fi

    echo "[$FILENAME] Display: ${DISP_W}x${STORED_H} -> $RES_LABEL"

    # ================================================================
    # JAPANESE AUDIO — find first/default jpn track, map by index
    # ================================================================
    JPN_ABS=""
    JPN_REL=0

    ffprobe -v error -select_streams a \
        -show_entries stream=index,disposition=default:stream_tags=language \
        -of csv=p=0 "$FILENAME" > "$TMPFILE" 2>&1

    # First pass: default-flagged jpn audio
    while IFS= read -r LINE; do
        if [[ -z "$JPN_ABS" ]]; then
            if echo "$LINE" | grep -qi "jpn"; then
                if echo "$LINE" | grep -qP ",1,|,1$"; then
                    JPN_ABS=$(echo "$LINE" | cut -d',' -f1)
                fi
            fi
        fi
    done < "$TMPFILE"

    # Second pass: first jpn audio track
    if [[ -z "$JPN_ABS" ]]; then
        while IFS= read -r LINE; do
            if [[ -z "$JPN_ABS" ]]; then
                if echo "$LINE" | grep -qi "jpn"; then
                    JPN_ABS=$(echo "$LINE" | cut -d',' -f1)
                fi
            fi
        done < "$TMPFILE"
    fi

    # Convert absolute stream index to audio-relative index
    if [[ -n "$JPN_ABS" ]]; then
        JPN_REL=0
        ffprobe -v error -select_streams a \
            -show_entries stream=index -of csv=p=0 "$FILENAME" > "$TMPFILE" 2>&1
        JPN_FOUND=""
        while IFS= read -r LINE; do
            IDX=$(echo "$LINE" | cut -d',' -f1)
            if [[ "$IDX" == "$JPN_ABS" ]]; then
                JPN_FOUND=1
                break
            fi
            (( JPN_REL++ )) || true
        done < "$TMPFILE"
        echo "[$FILENAME] Japanese audio: stream $JPN_ABS (audio index $JPN_REL)"
        JPN_MAP="0:a:$JPN_REL"
    else
        echo "[$FILENAME] WARNING: No Japanese audio found, using first audio track."
        JPN_MAP="0:a:0"
    fi

    # ================================================================
    # ENCODE
    # ================================================================
    TITLE="$BASENAME"
    echo "[$FILENAME] Starting encode ($ENC_LABEL)..."

    # Escape the filename for the subtitles filter (colons and special chars)
    ESCAPED_PATH=$(printf '%s' "$FILENAME" | sed "s/'/'\\\\''/g; s/:/\\\\:/g")

    set +e
    if [[ "$USE_VAAPI" -eq 1 ]]; then
        # VA-API: CPU-side filters run first. format=yuv420p converts Hi10P
        # sources to 8-bit, then format=nv12,hwupload hands frames to the GPU.
        # shellcheck disable=SC2086
        ffmpeg -y \
            -vaapi_device "$VAAPI_DEVICE" \
            -i "$FILENAME" \
            -vf "${SCALE_FILTER},subtitles='${ESCAPED_PATH}':si=${SUB_INDEX},unsharp=${SHARPEN},format=yuv420p,format=nv12,hwupload" \
            -c:v "$ENCODER" \
            $ENC_PARAMS \
            -tag:v hvc1 \
            -movflags +faststart \
            -c:a aac \
            -b:a "$AUDIO_BITRATE" \
            -ac 2 \
            -map 0:v:0 \
            -map "$JPN_MAP" \
            -map_chapters 0 \
            -metadata title="$TITLE" \
            "${OUTPUT_DIR}/${BASENAME}.mp4"
    else
        # AMF, NVENC, libx265: format=yuv420p in the filter chain handles
        # Hi10P sources by converting to 8-bit before encoding.
        # shellcheck disable=SC2086
        ffmpeg -y \
            -i "$FILENAME" \
            -vf "${SCALE_FILTER},subtitles='${ESCAPED_PATH}':si=${SUB_INDEX},unsharp=${SHARPEN},format=yuv420p" \
            -c:v "$ENCODER" \
            $ENC_PARAMS \
            -tag:v hvc1 \
            -movflags +faststart \
            -c:a aac \
            -b:a "$AUDIO_BITRATE" \
            -ac 2 \
            -map 0:v:0 \
            -map "$JPN_MAP" \
            -map_chapters 0 \
            -metadata title="$TITLE" \
            "${OUTPUT_DIR}/${BASENAME}.mp4"
    fi
    EXIT_CODE=$?
    set -e

    if [[ $EXIT_CODE -eq 0 ]]; then
        echo "[$FILENAME] Done > ${OUTPUT_DIR}/${BASENAME}.mp4"
        (( ENCODED++ )) || true
    else
        echo "[$FILENAME] FAILED - check output above for errors"
        (( FAILED++ )) || true
    fi

    echo
done

echo "============================================================"
echo " Finished!  Encoder: $ENC_LABEL"
echo " Encoded: $ENCODED   Failed: $FAILED   Skipped: $SKIPPED"
echo "============================================================"
