#!/bin/bash

set -e

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

INPUT="$1"
DELAY_MS="$2"
DELAY_SEC=$(awk "BEGIN {print $DELAY_MS / 1000}")
BASENAME="${INPUT%.*}"
OUTPUT="${BASENAME}_subdelay.mkv"
TMPDIR=$(mktemp -d)

echo "📦 Extracting subtitle stream info from $INPUT..."

# Get subtitle stream index, codec name, language, and title
mapfile -t SUB_TRACKS < <(ffprobe -v error -select_streams s -show_entries stream=index,codec_name -show_entries stream_tags=title,language -of csv=p=0 "$INPUT")

if [ ${#SUB_TRACKS[@]} -eq 0 ]; then
  echo "❌ No subtitle tracks found."
  rm -r "$TMPDIR"
  exit 1
fi

echo "🔍 Found ${#SUB_TRACKS[@]} subtitle track(s)"

# Extract the subtitle tracks
EXTRACT_ARGS=()
for i in "${!SUB_TRACKS[@]}"; do
  TRACK_LINE="${SUB_TRACKS[$i]}"
  IDX=$(cut -d',' -f1 <<< "$TRACK_LINE")
  CODEC=$(cut -d',' -f2 <<< "$TRACK_LINE")
  LANG=$(cut -d',' -f3 <<< "$TRACK_LINE")
  TITLE=$(cut -d',' -f4 <<< "$TRACK_LINE")
  
  case "$CODEC" in
    subrip) EXT="srt" ;;
    ass) EXT="ass" ;;
    ssa) EXT="ssa" ;;
    *)
      echo "⚠️ Unsupported subtitle codec: $CODEC (track $IDX). Skipping."
      continue
      ;;
  esac

  OUTFILE="$TMPDIR/sub${i}.${EXT}"
  EXTRACT_ARGS+=("$IDX:$OUTFILE")
  LANGS+=("$LANG")
  TITLES+=("$TITLE")
done

if [ ${#EXTRACT_ARGS[@]} -eq 0 ]; then
  echo "❌ No supported subtitle tracks found (only .srt/.ass/.ssa supported)."
  rm -r "$TMPDIR"
  exit 1
fi

echo "📤 Extracting subtitles..."
mkvextract tracks "$INPUT" "${EXTRACT_ARGS[@]}"

# Apply the delay to each subtitle file
for i in "${!EXTRACT_ARGS[@]}"; do
  EXT="${EXTRACT_ARGS[$i]##*.}"  # Get the file extension (e.g. ass, srt)
  IN="$TMPDIR/sub${i}.${EXT}"
  OUT="$TMPDIR/sub${i}_delayed.${EXT}"

  echo "⏱️ Delaying subtitle track $((i+1)) (${EXT}) by ${DELAY_MS}ms..."
  
  # Apply delay to the subtitle using FFmpeg
  ffmpeg -y -itsoffset "$DELAY_SEC" -i "$IN" -c copy "$OUT"
done

# Build the mkvmerge command to remux with the delayed subtitles
CMD=(mkvmerge -o "$OUTPUT" -S "$INPUT")

# Correct mapping for subtitle tracks with language and title flags
for i in "${!EXTRACT_ARGS[@]}"; do
  EXT="${EXTRACT_ARGS[$i]##*.}"
  REMUX_FILE="$TMPDIR/sub${i}_delayed.${EXT}"
  
  # Add the language and title metadata for each subtitle track
  CMD+=("--language" "0:${LANGS[$i]}" "--track-name" "0:${TITLES[$i]}" "$REMUX_FILE")
done

# Run the remuxing command
echo "🛠️ Remuxing with delayed subtitles..."
echo ${CMD[@]}
"${CMD[@]}"

echo "✅ Done! Output written to: $OUTPUT"
rm -r "$TMPDIR"