#!/bin/bash
set -euo pipefail
# -------------------------------
# Check input
# -------------------------------
if [ $# -lt 1 ]; then
    echo "Usage: $0 <input-file.mkv>"
    exit 1
fi
INPUT="$1"
BASENAME=$(basename "$INPUT" | sed 's/\.[^.]*$//')
OUTDIR="${BASENAME}-Transcode"
echo "========================================"
echo "Processing: $INPUT"
echo "Output: $OUTDIR"
echo "========================================"
mkdir -p "$OUTDIR/sub-titles"
mkdir -p "$OUTDIR/video"
mkdir -p "$OUTDIR/audio"
# -------------------------------
# 0️⃣ Generate mediainfo.json
# -------------------------------
echo "[0/5] Generating mediainfo.json"
ffprobe -v quiet -print_format json -show_format -show_streams "$INPUT" > "$OUTDIR/mediainfo.json"
# -------------------------------
# 1️⃣ Extract subtitles
# -------------------------------
echo "[1/5] Extracting subtitles"
IDX=0
for i in $(ffprobe -v error -select_streams s -show_entries stream=index,codec_name:stream_tags=language -of csv=p=0 "$INPUT"); do
    INDEX=$(echo $i | cut -d',' -f1)
    CODEC=$(echo $i | cut -d',' -f2)
    LANG=$(echo $i | cut -d',' -f3)
    LANG=${LANG:-und}
    if [[ "$CODEC" == "subrip" || "$CODEC" == "ass" || "$CODEC" == "ssa" ]]; then
        # Text-based → keep as .srt
        FNAME="$OUTDIR/sub-titles/${IDX}_${LANG}.srt"
        ffmpeg -y -i "$INPUT" -map 0:$INDEX -c:s srt "$FNAME"
        echo "  Extracted text subtitle: $FNAME"
    elif [[ "$CODEC" == "dvd_subtitle" ]]; then
        # Image-based → keep as .mkv
        FNAME_BASE="$OUTDIR/sub-titles/${IDX}_${LANG}"
        ffmpeg -y -i "$INPUT" -map 0:$INDEX -c:s dvd_subtitle "$FNAME_BASE.mkv"
        echo "  Extracted image subtitle: $FNAME_BASE.mkv"
    else
        echo "  Skipping unknown subtitle codec: $CODEC"
    fi
    IDX=$((IDX+1))
done
# -------------------------------
# 2️⃣ Encode video streams (GPU) with 10min chunks
# -------------------------------
echo "[2/5] Encoding video streams"
# Define profiles (resolution: bitrate)
PROFILES=("1080p:2000" "720p:1200")

# Define aspect ratios for 1080p (format: "aspect_ratio:width:height")
declare -A RESOLUTIONS_1080P
RESOLUTIONS_1080P["0.80"]="864:1080"
RESOLUTIONS_1080P["1.25"]="1350:1080"
RESOLUTIONS_1080P["1.33"]="1440:1080"
RESOLUTIONS_1080P["1.66"]="1800:1080"
RESOLUTIONS_1080P["1.78"]="1920:1080"
RESOLUTIONS_1080P["1.85"]="1920:1038"
RESOLUTIONS_1080P["1.90"]="1920:1010"
RESOLUTIONS_1080P["2.00"]="1920:960"
RESOLUTIONS_1080P["2.35"]="1920:816"
RESOLUTIONS_1080P["2.37"]="1920:810"
RESOLUTIONS_1080P["2.39"]="1920:802"
RESOLUTIONS_1080P["2.40"]="1920:800"
RESOLUTIONS_1080P["2.44"]="1920:786"

# Define aspect ratios for 720p (format: "aspect_ratio:width:height")
declare -A RESOLUTIONS_720P
RESOLUTIONS_720P["0.80"]="576:720"
RESOLUTIONS_720P["1.25"]="900:720"
RESOLUTIONS_720P["1.33"]="960:720"
RESOLUTIONS_720P["1.66"]="1200:720"
RESOLUTIONS_720P["1.78"]="1280:720"
RESOLUTIONS_720P["1.85"]="1280:692"
RESOLUTIONS_720P["1.90"]="1280:674"
RESOLUTIONS_720P["2.00"]="1280:640"
RESOLUTIONS_720P["2.35"]="1280:544"
RESOLUTIONS_720P["2.37"]="1280:540"
RESOLUTIONS_720P["2.39"]="1280:536"
RESOLUTIONS_720P["2.40"]="1280:532"
RESOLUTIONS_720P["2.44"]="1280:524"

# -------------------------------
# Iterate through each profile
# -------------------------------
for profile in "${PROFILES[@]}"; do
    RES=${profile%%:*}
    BITRATE=${profile##*:}
    mkdir -p "$OUTDIR/video/$RES"
    
    # Get the input video resolution (width and height)
    INPUT_WIDTH=$(ffprobe -v error -select_streams v:0 -show_entries stream=width -of default=noprint_wrappers=1:nokey=1 "$INPUT")
    INPUT_HEIGHT=$(ffprobe -v error -select_streams v:0 -show_entries stream=height -of default=noprint_wrappers=1:nokey=1 "$INPUT")
    
    # Calculate the input aspect ratio
    ASPECT_RATIO=$(echo "scale=2; $INPUT_WIDTH / $INPUT_HEIGHT" | bc)
    
    # Determine the resolution based on the aspect ratio and target resolution (1080p or 720p)
    if [ "$RES" == "1080p" ]; then
        if [[ -v RESOLUTIONS_1080P["$ASPECT_RATIO"] ]]; then
            RESOLUTION="${RESOLUTIONS_1080P[$ASPECT_RATIO]}"
            WIDTH="${RESOLUTION%%:*}"
            HEIGHT="${RESOLUTION##*:}"
            echo "Aspect ratio $ASPECT_RATIO found for 1080p. Using resolution $WIDTH x $HEIGHT."
        else
            echo "Unknown aspect ratio $ASPECT_RATIO for 1080p. Falling back to 16:9."
            WIDTH=1920
            HEIGHT=1080
        fi
    elif [ "$RES" == "720p" ]; then
        if [[ -v RESOLUTIONS_720P["$ASPECT_RATIO"] ]]; then
            RESOLUTION="${RESOLUTIONS_720P[$ASPECT_RATIO]}"
            WIDTH="${RESOLUTION%%:*}"
            HEIGHT="${RESOLUTION##*:}"
            echo "Aspect ratio $ASPECT_RATIO found for 720p. Using resolution $WIDTH x $HEIGHT."
        else
            echo "Unknown aspect ratio $ASPECT_RATIO for 720p. Falling back to 16:9."
            WIDTH=1280
            HEIGHT=720
        fi
    fi
    
    # Run FFmpeg with the correct scaling options (video only, no audio)
    ffmpeg -y -hwaccel cuda -i "$INPUT" \
        -map 0:v:0 -c:v h264_nvenc -b:v "${BITRATE}k" \
        -vf "scale=$WIDTH:$HEIGHT,setsar=1" \
        -profile:v high -pix_fmt yuv420p \
        -an \
        -f hls -hls_time 600 -hls_list_size 0 -hls_flags delete_segments+append_list \
        "$OUTDIR/video/$RES/stream.m3u8"
    
    echo "  $RES @ $BITRATE kbps encoded (Aspect Ratio: $ASPECT_RATIO)"
done

# -------------------------------
# 3️⃣ Encode audio tracks separately (10min chunks)
# -------------------------------
echo "[3/5] Encoding audio streams"
AUDIO_STREAMS=$(jq -r '.streams[] | select(.codec_type=="audio") | @base64' "$OUTDIR/mediainfo.json")
IDX=0
for s in $AUDIO_STREAMS; do
    _jq() { echo "$s" | base64 --decode | jq -r "$1"; }
    LANG=$(_jq '.tags.language // "und"')
    TITLE=$(_jq '.tags.title // "Track"')
    STREAM_INDEX=$(_jq '.index')
    AUDIO_OUT="$OUTDIR/audio/${IDX}_${LANG}"
    mkdir -p "$AUDIO_OUT"
    
    ffmpeg -y -i "$INPUT" -map 0:$STREAM_INDEX -c:a aac -b:a 128k \
        -vn \
        -f hls -hls_time 600 -hls_list_size 0 -hls_flags delete_segments+append_list \
        "$AUDIO_OUT/stream.m3u8"
    
    echo "  Extracted audio track $IDX (${LANG}) - ${TITLE}"
    IDX=$((IDX+1))
done

# -------------------------------
# 4️⃣ Generate master HLS playlist (JW Player compatible)
# -------------------------------
echo "[4/5] Creating master playlist"
MASTER="$OUTDIR/playlist.m3u8"
rm -f "$MASTER"
echo "#EXTM3U" >> "$MASTER"
echo "#EXT-X-VERSION:3" >> "$MASTER"
echo "" >> "$MASTER"

# Add audio tracks first (with proper GROUP-ID)
AUDIO_STREAMS=$(jq -r '.streams[] | select(.codec_type=="audio") | @base64' "$OUTDIR/mediainfo.json")
IDX=0
FIRST_AUDIO=true
for s in $AUDIO_STREAMS; do
    _jq() { echo "$s" | base64 --decode | jq -r "$1"; }
    LANG=$(_jq '.tags.language // "und"')
    TITLE=$(_jq '.tags.title // "Track"')
    
    # Clean title for display
    DISPLAY_NAME="${TITLE}"
    
    # Differentiate between audio types based on title or language, use better naming for tracks
    if [[ "$TITLE" == *"Audio Description"* ]]; then
        GROUP_ID="audio-description"
        DISPLAY_NAME="Audio Description (${LANG})"
    elif [[ "$TITLE" == *"Dialogue Boost"* ]]; then
        GROUP_ID="dialogue-boost"
        DISPLAY_NAME="Dialogue Boost (${LANG})"
    else
        GROUP_ID="audio"
        DISPLAY_NAME="${LANG} (${TITLE})"
    fi
    
    # First audio track is default
    if [ "$FIRST_AUDIO" = true ]; then
        echo "#EXT-X-MEDIA:TYPE=AUDIO,GROUP-ID=\"${GROUP_ID}\",NAME=\"${DISPLAY_NAME}\",LANGUAGE=\"${LANG}\",DEFAULT=YES,AUTOSELECT=YES,URI=\"audio/${IDX}_${LANG}/stream.m3u8\"" >> "$MASTER"
        FIRST_AUDIO=false
    else
        echo "#EXT-X-MEDIA:TYPE=AUDIO,GROUP-ID=\"${GROUP_ID}\",NAME=\"${DISPLAY_NAME}\",LANGUAGE=\"${LANG}\",DEFAULT=NO,AUTOSELECT=YES,URI=\"audio/${IDX}_${LANG}/stream.m3u8\"" >> "$MASTER"
    fi
    IDX=$((IDX+1))
done

echo "" >> "$MASTER"

# Add subtitles (SRT and MKV)
for s in "$OUTDIR/sub-titles"/*; do
    if [ -f "$s" ]; then
        SUB_NAME=$(basename "$s")
        EXT="${SUB_NAME##*.}"
        
        # Only add text-based subtitles (SRT) to HLS
        # MKV files are kept for external loading if needed
        if [[ "$EXT" == "srt" ]]; then
            # Extract language from filename (format: IDX_LANG.ext)
            SUB_BASE="${SUB_NAME%.*}"
            SUB_LANG="${SUB_BASE##*_}"
            DISPLAY_LANG=$(echo "$SUB_LANG" | tr '[:lower:]' '[:upper:]')
            
            echo "#EXT-X-MEDIA:TYPE=SUBTITLES,GROUP-ID=\"subs\",NAME=\"${DISPLAY_LANG}\",LANGUAGE=\"${SUB_LANG}\",DEFAULT=NO,AUTOSELECT=YES,FORCED=NO,URI=\"sub-titles/${SUB_NAME}\"" >> "$MASTER"
        fi
    fi
done

echo "" >> "$MASTER"

# Add video variants with audio and subtitle groups
for profile in "${PROFILES[@]}"; do
    RES=${profile%%:*}
    BITRATE=${profile##*:}
    
    # Get resolution from the first segment to be accurate
    if [ -f "$OUTDIR/video/$RES/stream.m3u8" ]; then
        # Try to get actual resolution from ffprobe
        FIRST_SEGMENT=$(grep -v '^#' "$OUTDIR/video/$RES/stream.m3u8" | head -1)
        if [ -n "$FIRST_SEGMENT" ]; then
            SEGMENT_PATH="$OUTDIR/video/$RES/$FIRST_SEGMENT"
            if [ -f "$SEGMENT_PATH" ]; then
                VID_WIDTH=$(ffprobe -v error -select_streams v:0 -show_entries stream=width -of default=noprint_wrappers=1:nokey=1 "$SEGMENT_PATH")
                VID_HEIGHT=$(ffprobe -v error -select_streams v:0 -show_entries stream=height -of default=noprint_wrappers=1:nokey=1 "$SEGMENT_PATH")
            fi
        fi
        
        # Fallback to standard resolutions if detection failed
        if [ -z "$VID_WIDTH" ] || [ -z "$VID_HEIGHT" ]; then
            if [ "$RES" == "1080p" ]; then
                VID_WIDTH=1920
                VID_HEIGHT=1080
            else
                VID_WIDTH=1280
                VID_HEIGHT=720
            fi
        fi
        
        # Check if we have subtitles
        HAS_SUBS=$(ls -1 "$OUTDIR/sub-titles"/*.srt 2>/dev/null | wc -l)
        
        if [ "$HAS_SUBS" -gt 0 ]; then
            echo "#EXT-X-STREAM-INF:BANDWIDTH=$((BITRATE*1024)),RESOLUTION=${VID_WIDTH}x${VID_HEIGHT},CODECS=\"avc1.640028,mp4a.40.2\",AUDIO=\"audio\",SUBTITLES=\"subs\"" >> "$MASTER"
        else
            echo "#EXT-X-STREAM-INF:BANDWIDTH=$((BITRATE*1024)),RESOLUTION=${VID_WIDTH}x${VID_HEIGHT},CODECS=\"avc1.640028,mp4a.40.2\",AUDIO=\"audio\"" >> "$MASTER"
        fi
        echo "video/$RES/stream.m3u8" >> "$MASTER"
    fi
done

echo "[5/5] Transcoding complete!"
echo "Master playlist: $MASTER"
echo "========================================"
