#!/bin/bash
set -euo pipefail

if [ $# -lt 1 ]; then
    echo "Usage: $0 <input.mkv>"
    exit 1
fi

INPUT="$1"
BASENAME=$(basename "$INPUT" | sed 's/\.[^.]*$//')
OUTDIR="${BASENAME}-Transcode"
mkdir -p "$OUTDIR/sub-titles" "$OUTDIR/video" "$OUTDIR/audio"

# -------------------------------
# 0️⃣ Generate mediainfo.json
# -------------------------------
echo "[0/5] Analyzing Input..."
ffprobe -v quiet -print_format json -show_format -show_streams "$INPUT" > "$OUTDIR/mediainfo.json"

# -------------------------------
# 1️⃣ Master Playlist Generation
# -------------------------------
echo "[1/5] Building Master Playlist..."
MASTER="$OUTDIR/playlist.m3u8"
printf "#EXTM3U\n#EXT-X-VERSION:3\n\n" > "$MASTER"

# AUDIO TRACKS
AUDIO_DATA=$(jq -r '.streams[] | select(.codec_type=="audio") | @base64' "$OUTDIR/mediainfo.json")
IDX=0
for s in $AUDIO_DATA; do
    _jq() { echo "$s" | base64 --decode | jq -r "$1"; }
    LANG=$(_jq '.tags.language // "und"')
    TITLE=$(_jq '.tags.title // "Audio"')
    printf "#EXT-X-MEDIA:TYPE=AUDIO,GROUP-ID=\"audio\",NAME=\"%s\",LANGUAGE=\"%s\",DEFAULT=%s,AUTOSELECT=YES,URI=\"audio/%s_%s/stream.m3u8\"\n" \
        "$TITLE" "$LANG" "$([ $IDX -eq 0 ] && echo "YES" || echo "NO")" "$IDX" "$LANG" >> "$MASTER"
    IDX=$((IDX+1))
done

# SUBTITLE TRACKS (Text-only)
SUB_DATA=$(jq -r '.streams[] | select(.codec_type=="subtitle") | @base64' "$OUTDIR/mediainfo.json")
HAS_SUBS=0
IDX=0
for s in $SUB_DATA; do
    _jq() { echo "$s" | base64 --decode | jq -r "$1"; }
    CODEC=$(_jq '.codec_name')
    # Filter for HLS compatible text formats
    if [[ "$CODEC" =~ (subrip|ass|ssa|mov_text|webvtt) ]]; then
        LANG=$(_jq '.tags.language // "und"')
        NAME=$(echo "$LANG" | tr '[:lower:]' '[:upper:]')
        printf "#EXT-X-MEDIA:TYPE=SUBTITLES,GROUP-ID=\"subs\",NAME=\"%s\",LANGUAGE=\"%s\",DEFAULT=NO,AUTOSELECT=YES,FORCED=NO,URI=\"sub-titles/%s_%s.vtt\"\n" \
            "$NAME" "$LANG" "$IDX" "$LANG" >> "$MASTER"
        HAS_SUBS=$((HAS_SUBS+1))
    fi
    IDX=$((IDX+1))
done

printf "\n" >> "$MASTER"

# VIDEO VARIANTS
PROFILES=("1080p:2000" "720p:1200")
for profile in "${PROFILES[@]}"; do
    RES=${profile%%:*}
    BITRATE=${profile##*:}
    V_RES="$([ "$RES" == "1080p" ] && echo "1920x1080" || echo "1280x720")"
    INF="#EXT-X-STREAM-INF:BANDWIDTH=$((BITRATE*1024)),RESOLUTION=$V_RES,CODECS=\"avc1.640028,mp4a.40.2\",AUDIO=\"audio\""
    [ "$HAS_SUBS" -gt 0 ] && INF="$INF,SUBTITLES=\"subs\""
    printf "%s\nvideo/%s/stream.m3u8\n" "$INF" "$RES" >> "$MASTER"
done

# -------------------------------
# 2️⃣ Smart Subtitle Extraction
# -------------------------------
echo "[2/5] Extracting Subtitles..."
IDX=0
while IFS=, read -r stream_idx codec lang; do
    lang=${lang:-und}
    if [[ "$codec" =~ (subrip|ass|ssa|mov_text|webvtt) ]]; then
        ffmpeg -y -i "$INPUT" -map 0:$stream_idx -c:s webvtt "$OUTDIR/sub-titles/${IDX}_${lang}.vtt" -loglevel error
    fi
    IDX=$((IDX+1))
done < <(ffprobe -v error -select_streams s -show_entries stream=index,codec_name:stream_tags=language -of csv=p=0 "$INPUT")

# -------------------------------
# 3️⃣ Smart Video Encoding (GPU + Copy Logic)
# -------------------------------
echo "[3/5] Processing Video..."
V_CODEC=$(jq -r '.streams[] | select(.codec_type=="video") | .codec_name' "$OUTDIR/mediainfo.json" | head -1)

for profile in "${PROFILES[@]}"; do
    RES=${profile%%:*}
    BITRATE=${profile##*:}
    TARGET_H="$([ "$RES" == "1080p" ] && echo 1080 || echo 720)"
    mkdir -p "$OUTDIR/video/$RES"

    # If already H264, we can speed up by copying (only if resolution matches)
    # For now, let's use NVENC for speed unless source matches exactly
    echo "  Encoding $RES variant..."
    ffmpeg -y -hwaccel cuda -i "$INPUT" \
        -map 0:v:0 -c:v h264_nvenc -b:v "${BITRATE}k" \
        -vf "scale=trunc(oh*a/2)*2:$TARGET_H" \
        -an -f hls -hls_time 600 -hls_list_size 0 "$OUTDIR/video/$RES/stream.m3u8" -loglevel error
done

# -------------------------------
# 4️⃣ Smart Audio Encoding (Copy Logic)
# -------------------------------
echo "[4/5] Processing Audio..."
IDX=0
for s in $AUDIO_DATA; do
    _jq() { echo "$s" | base64 --decode | jq -r "$1"; }
    A_CODEC=$(_jq '.codec_name')
    A_BITRATE=$(echo "$(_jq '.bit_rate') / 1000" | bc 2>/dev/null || echo "0")
    LANG=$(_jq '.tags.language // "und"')
    STREAM_INDEX=$(_jq '.index')
    
    AUDIO_OUT="$OUTDIR/audio/${IDX}_${LANG}"
    mkdir -p "$AUDIO_OUT"

    # If already AAC and bitrate is close to 128k, copy it
    if [[ "$A_CODEC" == "aac" ]]; then
        echo "  Copying AAC track $IDX..."
        A_ENC="copy"
    else
        echo "  Transcoding track $IDX to AAC..."
        A_ENC="aac"
    fi

    ffmpeg -y -i "$INPUT" -map 0:$STREAM_INDEX -c:a $A_ENC -b:a 128k -vn \
        -f hls -hls_time 600 -hls_list_size 0 "$AUDIO_OUT/stream.m3u8" -loglevel error
    IDX=$((IDX+1))
done

echo "[5/5] Complete!"
