#!/usr/bin/env bash
# Cut dead air, burn captions, export h264 mp4 into exports/.
#
# Docs:
#   https://ffmpeg.org/ffmpeg-filters.html#silencedetect
#   https://ffmpeg.org/ffmpeg-filters.html#subtitles
#   https://github.com/SYSTRAN/faster-whisper
#
# silencedetect defaults are noise=-60dB and duration=2. Talking-head room
# tone sits above -60dB, so this script uses -30dB and 0.8s unless overridden.
# SILENCE_PAD (default 0.15s) is a script choice so a cut keeps the word edge.
# Each removed span is appended to cuts.txt in source time, so a jump that
# was already in the footage is absent from that log.
# --end-on-sentence (default on) trims the export to the last word that ends
# in . ? or !, rewrites the SRT from the words that remain, and saves
# exports/<name>.srt plus exports/<name>.words.json.
#
# On Linux, whisper runs only after MemAvailable is at least 3GB (skipped when
# /proc/meminfo is missing, e.g. macOS), and holds flock /tmp/epic-heavy-job.lock
# when flock exists (skipped otherwise). Portable bash 3.2+ (macOS default).
# whisper-srt.py waits again before it loads
# the model, so the lock holder starts only after that check.

set -euo pipefail

export LC_ALL=C

SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)

SILENCE_NOISE=${SILENCE_NOISE:--30dB}
SILENCE_DURATION=${SILENCE_DURATION:-0.8}
SILENCE_PAD=${SILENCE_PAD:-0.15}
WHISPER_MODEL=${WHISPER_MODEL:-small.en}

usage() {
  cat >&2 <<'EOF'
usage: ai-editor-edit.sh [--end-on-sentence|--no-end-on-sentence] <raw-video>

Cut dead air, burn one-line captions, write h264 mp4 to exports/<name>.mp4.
The same folder gets exports/<name>.srt and exports/<name>.words.json.
Run from the editor project (raw-footage/, references/, exports/), or pass a
file that already lives under raw-footage/.

--end-on-sentence is the default. It trims the export to the last word that
ends in . ? or !, and rewrites the SRT so the last cue ends on that word.
--no-end-on-sentence keeps the take through the last word.

Writes cuts.txt in the project: source timestamps this script removed.
A jump in the picture that is missing from cuts.txt was already in the source.

Env:
  SILENCE_NOISE      default -30dB   (silencedetect noise; filter default is -60dB)
  SILENCE_DURATION   default 0.8     (seconds; filter default is 2)
  SILENCE_PAD        default 0.15    (seconds of silence kept on each side of a cut)
  WHISPER_MODEL      default small.en (or base.en)
EOF
}

validate_env() {
  case "$WHISPER_MODEL" in
    small.en|base.en) ;;
    *)
      echo "WHISPER_MODEL must be small.en or base.en" >&2
      exit 1
      ;;
  esac
  if [[ ! "$SILENCE_NOISE" =~ ^-?[0-9]+([.][0-9]+)?(dB)?$ ]]; then
    echo "SILENCE_NOISE must look like -30dB or 0.001" >&2
    exit 1
  fi
  if [[ ! "$SILENCE_DURATION" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
    echo "SILENCE_DURATION must be seconds, such as 0.8" >&2
    exit 1
  fi
  if [[ ! "$SILENCE_PAD" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
    echo "SILENCE_PAD must be seconds, such as 0.15" >&2
    exit 1
  fi
}

require_tools() {
  local cmd
  local missing=()
  for cmd in ffmpeg ffprobe python3 awk; do
    if ! command -v "$cmd" >/dev/null 2>&1; then
      missing+=("$cmd")
    fi
  done
  if [[ ${#missing[@]} -gt 0 ]]; then
    echo "missing command: ${missing[*]}" >&2
    exit 1
  fi
  local filters name
  filters=$(ffmpeg -hide_banner -filters 2>/dev/null || true)
  for name in silencedetect subtitles concat; do
    if ! awk -v n="$name" '$2 == n { found = 1 } END { exit !found }' <<< "$filters"; then
      echo "ffmpeg is missing the ${name} filter. subtitles needs a build with libass." >&2
      echo "Docs: https://ffmpeg.org/ffmpeg-filters.html#${name}" >&2
      exit 1
    fi
  done
}

mem_available_kb() {
  local kb
  kb=$(awk '/^MemAvailable:/ { print $2; found = 1 } END { if (!found) exit 1 }' /proc/meminfo) || return 1
  if [[ -z "$kb" ]]; then
    return 1
  fi
  printf '%s\n' "$kb"
}

wait_for_free_memory() {
  local need avail
  # 3 GiB, in KiB. That is past 3GB, so whisper starts only with 3GB or more free.
  need=$((3 * 1024 * 1024))
  if [[ ! -e /proc/meminfo ]]; then
    echo "no /proc/meminfo (not Linux); skipping the free-memory wait" >&2
    return 0
  fi
  while true; do
    avail=$(mem_available_kb) || exit 1
    if [[ "$avail" -ge "$need" ]]; then
      echo "free memory ok: MemAvailable ${avail} kB (need ${need} kB)" >&2
      return 0
    fi
    echo "waiting for >=3GB free memory (MemAvailable ${avail} kB, need ${need} kB)" >&2
    sleep 10
  done
}

parse_silences() {
  awk '
    {
      gsub(/\r/, "")
    }
    /silence_start:/ {
      start = ""
      for (i = 1; i <= NF; i++) if ($i == "silence_start:") start = $(i + 1)
    }
    /silence_end:/ {
      end = ""
      for (i = 1; i <= NF; i++) if ($i == "silence_end:") end = $(i + 1)
      if (start != "") {
        print start, end
        start = ""
      }
    }
    END {
      if (start != "") print start, "EOF"
    }
  '
}

# Read "start end" silence rows on stdin. Print speech ranges to keep.
keep_ranges() {
  local duration=$1
  local pad=$2
  awk -v dur="$duration" -v pad="$pad" '
    BEGIN { cursor = 0 }
    function emit(a, b) {
      if (a < 0) a = 0
      if (b > dur) b = dur
      if ((b - a) > (pad + 0.001))
        printf "%.6f %.6f\n", a, b
    }
    {
      s = $1 + 0
      if ($2 == "EOF") e = dur + 0
      else e = $2 + 0
      if (e < s) next
      s += pad
      e -= pad
      if (s < 0) s = 0
      if (e > dur) e = dur
      if (e <= s) next
      if (s > cursor) emit(cursor, s)
      if (e > cursor) cursor = e
    }
    END {
      if (dur > cursor) emit(cursor, dur)
    }
  '
}

ranges_are_full_file() {
  local keeps=$1
  local duration=$2
  local lines
  lines=$(grep -c . "$keeps" || true)
  if [[ "$lines" -ne 1 ]]; then
    return 1
  fi
  awk -v d="$duration" 'NR == 1 { exit !($1 <= 0.02 && (d - $2) <= 0.02) }' "$keeps"
}

write_cut_filter() {
  local keeps=$1
  local dest=$2
  local start end i=0 pairs=""
  : > "$dest"
  while read -r start end; do
    [[ -n "${start:-}" && -n "${end:-}" ]] || continue
    printf '[0:v]trim=start=%s:end=%s,setpts=PTS-STARTPTS[v%d];\n' "$start" "$end" "$i" >> "$dest"
    printf '[0:a]atrim=start=%s:end=%s,asetpts=PTS-STARTPTS[a%d];\n' "$start" "$end" "$i" >> "$dest"
    pairs+="[v${i}][a${i}]"
    i=$((i + 1))
  done < "$keeps"
  if [[ "$i" -eq 0 ]]; then
    echo "no speech ranges to cut" >&2
    return 1
  fi
  if [[ "$i" -eq 1 ]]; then
    sed 's/\[v0\]/[v]/; s/\[a0\]/[a]/' "$dest" > "${dest}.one"
    mv "${dest}.one" "$dest"
    return 0
  fi
  printf '%sconcat=n=%d:v=1:a=1[v][a]\n' "$pairs" "$i" >> "$dest"
}

# Log source spans this script removes. Jumps already in the footage are absent.
write_cuts_log() {
  local keeps=$1
  local duration=$2
  local dest=$3
  local label=$4
  awk -v dur="$duration" -v label="$label" '
    BEGIN {
      printf "# %s\n", label
      printf "# source timestamps removed by this script (start end seconds)\n"
      printf "# a jump in the picture that is absent here was already in the source\n"
    }
    NR == 1 {
      if ($1 > 0.02) printf "%.6f %.6f\n", 0.0, $1 + 0
      prev = $2 + 0
      next
    }
    {
      if (($1 + 0) - prev > 0.02) printf "%.6f %.6f\n", prev, $1 + 0
      prev = $2 + 0
    }
    END {
      if (NR == 0) exit
      if (dur - prev > 0.02) printf "%.6f %.6f\n", prev, dur
    }
  ' "$keeps" > "$dest"
}

# Write a cut file. Return 2 when the take has no dead air to remove.
cut_dead_air() {
  local input=$1
  local dest=$2
  local work=$3
  local duration silences keeps count
  # Called from `if`, where bash ignores set -e, so each step checks its status.
  duration=$(ffprobe -v error -show_entries format=duration -of default=nw=1:nk=1 "$input") || return 1
  duration=$(awk -v d="$duration" 'BEGIN { if (d + 0 <= 0) exit 1; printf "%.6f", d + 0 }') || {
    echo "could not read duration: $input" >&2
    return 1
  }
  silences=$work/silence.txt
  keeps=$work/keeps.txt
  # -map 0:a:0 decodes the audio and skips the picture.
  # Log lines are "silence_start:" / "silence_end:" per the filter docs.
  ffmpeg -hide_banner -nostats -loglevel info -i "$input" -map 0:a:0 \
    -af "silencedetect=noise=${SILENCE_NOISE}:duration=${SILENCE_DURATION}" \
    -f null - 2>&1 | parse_silences | sort -n -k1,1 > "$silences" || return 1
  keep_ranges "$duration" "$SILENCE_PAD" < "$silences" > "$keeps" || return 1
  if [[ ! -s "$keeps" ]]; then
    echo "no speech left after the dead-air cut: $input" >&2
    return 1
  fi
  write_cuts_log "$keeps" "$duration" "$work/cuts.txt" "$(basename "$input")" || return 1
  if ranges_are_full_file "$keeps" "$duration"; then
    echo "no dead air to cut" >&2
    return 2
  fi
  count=$(awk 'END { print NR }' "$keeps")
  echo "keeping ${count} speech range(s)" >&2
  write_cut_filter "$keeps" "$work/graph.txt" || return 1
  ffmpeg -hide_banner -nostats -loglevel error -y -i "$input" \
    -filter_complex_script "$work/graph.txt" \
    -map "[v]" -map "[a]" \
    -c:v libx264 -preset veryfast -crf 16 -pix_fmt yuv420p \
    -c:a aac -b:a 160k \
    "$dest" || return 1
}

burn_export() {
  local media=$1
  local srt=$2
  local dest=$3
  local end_time=${4:-}
  local srt_dir srt_name
  local -a limit=()
  # One encode thread and bitexact flags so the same picture is the same bytes.
  local -a exact=(
    -threads 1
    -fflags +bitexact
    -flags:v +bitexact
    -flags:a +bitexact
    -map_metadata -1
  )
  mkdir -p "$(dirname "$dest")"
  if [[ -n "$end_time" ]]; then
    limit=(-t "$end_time")
  fi
  if grep -q -- '-->' "$srt"; then
    srt_dir=$(cd "$(dirname "$srt")" && pwd)
    srt_name=$(basename "$srt")
    # Run from the SRT directory so the subtitles filename has no special chars.
    # WrapStyle=2 is libass "no wrap", so the burn stays one line.
    # force_style follows the quoting in the subtitles filter docs.
    (
      cd "$srt_dir"
      ffmpeg -hide_banner -nostats -loglevel error -y -i "$media" \
        ${exact[@]+"${exact[@]}"} \
        -vf "subtitles=${srt_name}:force_style='Fontname=DejaVu Sans,Fontsize=14,PrimaryColour=&H00FFFFFF,OutlineColour=&H00000000,BorderStyle=1,Outline=2,Alignment=2,MarginV=20,WrapStyle=2'" \
        ${limit[@]+"${limit[@]}"} \
        -c:v libx264 -preset medium -crf 20 -pix_fmt yuv420p \
        -c:a aac -b:a 160k -movflags +faststart \
        "$dest"
    )
  else
    echo "no caption cues; exporting h264 without a burn" >&2
    ffmpeg -hide_banner -nostats -loglevel error -y -i "$media" \
      ${exact[@]+"${exact[@]}"} \
      ${limit[@]+"${limit[@]}"} \
      -c:v libx264 -preset medium -crf 20 -pix_fmt yuv420p \
      -c:a aac -b:a 160k -movflags +faststart \
      "$dest"
  fi
}

project_root() {
  local input=$1
  local dir
  if [[ -d raw-footage || -d references || -d exports ]]; then
    pwd
    return
  fi
  dir=$(cd "$(dirname "$input")" && pwd)
  while [[ "$dir" != "/" ]]; do
    if [[ -d "$dir/raw-footage" || -d "$dir/references" || -d "$dir/exports" ]]; then
      printf '%s\n' "$dir"
      return
    fi
    if [[ "$(basename "$dir")" == "raw-footage" ]]; then
      dirname "$dir"
      return
    fi
    dir=$(dirname "$dir")
  done
  pwd
}

require_av() {
  local input=$1
  local video audio
  video=$(ffprobe -v error -select_streams v:0 -show_entries stream=codec_type -of csv=p=0 "$input" | head -n 1 || true)
  audio=$(ffprobe -v error -select_streams a:0 -show_entries stream=codec_type -of csv=p=0 "$input" | head -n 1 || true)
  if [[ "$video" != "video" ]]; then
    echo "no video stream: $input" >&2
    exit 1
  fi
  if [[ "$audio" != "audio" ]]; then
    echo "no audio stream: $input" >&2
    exit 1
  fi
}

main() {
  local input root stem dest tmp media srt status
  local end_on_sentence=1
  local end_time="" media_dur words_json
  local -a positional=()
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --end-on-sentence)
        end_on_sentence=1
        shift
        ;;
      --no-end-on-sentence)
        end_on_sentence=0
        shift
        ;;
      -h|--help)
        usage
        exit 2
        ;;
      --)
        shift
        positional+=("$@")
        break
        ;;
      -*)
        echo "unknown flag: $1" >&2
        usage
        exit 2
        ;;
      *)
        positional+=("$1")
        shift
        ;;
    esac
  done
  if [[ ${#positional[@]} -lt 1 ]]; then
    usage
    exit 2
  fi
  set -- ${positional[@]+"${positional[@]}"}
  validate_env
  require_tools
  if [[ -f "$1" ]]; then
    input=$(cd "$(dirname "$1")" && pwd)/$(basename "$1")
  elif [[ -f "raw-footage/$1" ]]; then
    input=$(cd "$(dirname "raw-footage/$1")" && pwd)/$(basename "$1")
  else
    echo "raw file not found: $1" >&2
    exit 1
  fi
  require_av "$input"
  root=$(project_root "$input")
  stem=$(basename "$input")
  stem=${stem%.*}
  if [[ -z "$stem" || "$stem" == "." ]]; then
    echo "could not name the export from $input" >&2
    exit 1
  fi
  mkdir -p "$root/exports"
  dest="$root/exports/${stem}.mp4"
  if [[ "$dest" == "$input" ]]; then
    echo "refusing to overwrite the raw file: $dest" >&2
    exit 1
  fi
  tmp=$(mktemp -d)
  # Bake the path into the trap. tmp is local to main, so a '$tmp' trap
  # fires after that local is gone and trips set -u.
  trap "rm -rf $(printf '%q' "$tmp")" EXIT
  media=$input
  if cut_dead_air "$input" "$tmp/cut.mp4" "$tmp"; then
    media=$tmp/cut.mp4
  else
    status=$?
    if [[ "$status" -ne 2 ]]; then
      exit "$status"
    fi
  fi
  cp "$tmp/cuts.txt" "$root/cuts.txt"
  echo "cuts log: $root/cuts.txt" >&2
  srt=$tmp/captions.srt
  words_json=$tmp/words.json
  # Wait, then caption under the shared heavy-job lock.
  wait_for_free_memory
  if command -v flock >/dev/null 2>&1; then
    flock /tmp/epic-heavy-job.lock python3 "$SCRIPT_DIR/whisper-srt.py" --model "$WHISPER_MODEL" "$media" "$srt"
  else
    echo "flock not found; running captions without the shared lock" >&2
    python3 "$SCRIPT_DIR/whisper-srt.py" --model "$WHISPER_MODEL" "$media" "$srt"
  fi
  if [[ "$end_on_sentence" -eq 1 ]]; then
    if [[ ! -f "$words_json" || ! -f "$srt" ]]; then
      echo "captions are missing; cannot trim to the last sentence" >&2
      exit 1
    fi
    # Rewrites the SRT and words.json, then prints the export end.
    end_time=$(python3 "$SCRIPT_DIR/whisper-srt.py" --apply-sentence-trim "$words_json" "$srt")
    if [[ -n "$end_time" ]]; then
      media_dur=$(ffprobe -v error -show_entries format=duration -of default=nw=1:nk=1 "$media")
      if awk -v end="$end_time" -v dur="$media_dur" 'BEGIN { exit !((dur - end) > 0.05 && end > 0.05) }'; then
        echo "trimming export to last sentence at ${end_time}s" >&2
      else
        end_time=""
      fi
    fi
  fi
  if [[ ! -f "$srt" || ! -f "$words_json" ]]; then
    echo "captions are missing" >&2
    exit 1
  fi
  cp "$srt" "$root/exports/${stem}.srt"
  cp "$words_json" "$root/exports/${stem}.words.json"
  echo "captions: $root/exports/${stem}.srt" >&2
  burn_export "$media" "$srt" "$dest" "$end_time"
  printf '%s\n' "$dest"
}

if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then
  main "$@"
fi