#!/usr/bin/env bash
# gpu-nvenc-test.sh —— 用新版 ffmpeg(带 -bf/-lookahead_level/-tune uhq) 跑 av1_nvenc 多档对比
# 用法： bash gpu-nvenc-test.sh <源文件> <输出目录> <ffmpeg路径>
set -uo pipefail

SRC="${1:?源文件}"
OUT="${2:?输出目录}"
FF="${3:-ffmpeg}"
SEGS=(900 3000 5200)
DUR=40
TSV="$OUT/nvenc.tsv"
mkdir -p "$OUT"
printf 'variant\tenc\tparams\tseg\tbytes\tkbps\tssim\tsecs\n' > "$TSV"

# 名字|编码参数（不含 -c:v 与 cq）
V=(
  "nv-p3-cq26|av1_nvenc|-preset p3 -tune hq -bf 6 -lookahead_level 2 -qp_cb_offset 12 -cq 26"
  "nv-p3-cq30|av1_nvenc|-preset p3 -tune hq -bf 6 -lookahead_level 2 -qp_cb_offset 12 -cq 30"
  "nv-p5uhq-cq26|av1_nvenc|-preset p5 -tune uhq -bf 6 -lookahead_level 2 -qp_cb_offset 12 -cq 26"
  "nv-p7-cq26|av1_nvenc|-preset p7 -tune hq -cq 26"
)

echo "═══ av1_nvenc 多档对比（新 ffmpeg）═══"
"$FF" -hide_banner -version | head -1 | sed 's/^/  /'

for item in "${V[@]}"; do
  name="${item%%|*}"; rest="${item#*|}"; enc="${rest%%|*}"; params="${rest#*|}"
  for s in "${SEGS[@]}"; do
    out="$OUT/av1-$name-$s.mp4"
    t0=$(date +%s.%N)
    if ! nice -n 19 "$FF" -hide_banner -nostdin -loglevel error -y -ss "$s" -t "$DUR" -i "$SRC" \
        -map 0:v:0 -map "0:a:0?" -c:a aac -b:a 128k \
        -c:v "$enc" $params -pix_fmt yuv420p -movflags +faststart "$out" 2>"$OUT/err-$name-$s.log"; then
      echo "  ❌ $name @ ${s}s 失败：$(tail -1 "$OUT/err-$name-$s.log")"; continue
    fi
    t1=$(date +%s.%N)
    secs=$(echo "$t1 - $t0" | bc)
    fps=$(echo "scale=1; $((DUR*30000/1001)) / $secs" | bc)
    slog="$OUT/ssim-$name-$s.log"
    ssim=$(nice -n 19 "$FF" -hide_banner -nostdin -loglevel error \
        -ss "$s" -t "$DUR" -i "$SRC" -i "$out" \
        -lavfi "[0:v][1:v]ssim=stats_file=$slog" -f null - 2>&1 | grep -o 'All:[0-9.]*' | head -1 | cut -d: -f2)
    b=$(stat -c%s "$out")
    kbps=$(( b * 8 / DUR / 1000 ))
    printf '%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n' "$name" "$enc" "$params" "$s" "$b" "$kbps" "${ssim:-?}" "$secs" >> "$TSV"
    echo "  $name @ ${s}s → ${secs%.*}s（$(printf %.1f "$(echo "$DUR/$secs"|bc -l)" )x 实时，${fps}fps）  $(du -h "$out"|cut -f1)  ${kbps}kbps  SSIM ${ssim:-?}"
  done
done

echo
echo "═══ 平均 ═══"
awk -F'\t' 'NR>1 && $7!="?" {n[$1]++; s[$1]+=$7; k[$1]+=$6; t[$1]+=$8}
  END{for(v in n) printf "  %-14s 平均SSIM %.5f  平均码率 %d kbps  平均耗时 %.1fs（%.2fx实时）\n", v, s[v]/n[v], k[v]/n[v], t[v]/n[v], 40/(t[v]/n[v])}' "$TSV" | sort