#!/usr/bin/env bash
# compare_snapshot.sh  (robust v2)
# Compare two MP4s by grabbing a snapshot at a given time and measuring SSIM/PSNR.
#
# Usage:
#   ./compare_snapshot.sh <file1.mp4> <file2.mp4> [time_seconds] [scale] [--verbose]
# Examples:
#   ./compare_snapshot.sh A.mp4 B.mp4
#   ./compare_snapshot.sh A.mp4 B.mp4 12 320:-1 --verbose
#
# Exit codes: 0 = ran ok (doesn't imply "very close"), 2 = usage error, 3 = ffmpeg failure

set -u  # keep -e/-o pipefail off to avoid silent exits from grep/awk mismatches

f1="${1:-}"; f2="${2:-}"
t="${3:-12}"
scale="${4:-320:-1}"
verbose=0
[[ "${5:-}" == "--verbose" ]] && verbose=1

if [[ -z "${f1}" || -z "${f2}" ]]; then
  echo "Usage: $0 <file1.mp4> <file2.mp4> [time_seconds] [scale] [--verbose]" >&2
  exit 2
fi
for f in "$f1" "$f2"; do
  [[ -f "$f" ]] || { echo "Not found: $f" >&2; exit 2; }
done

tmpdir="$(mktemp -d)"
trap 'rm -rf "$tmpdir"' EXIT

shot1="$tmpdir/shot1.png"
shot2="$tmpdir/shot2.png"
log="$tmpdir/metrics.log"

# Helper for verbose prints
vprint() { [[ $verbose -eq 1 ]] && echo "[DBG]" "$@" >&2; }

# 1) Extract frames (accurate seek: -ss after -i)
# Use -y to avoid “file exists” issues, and format to a stable pixel format.
cmd1=(ffmpeg -hide_banner -v error -y -i "$f1" -ss "$t" -frames:v 1 -vf "scale=$scale,format=yuv420p" "$shot1")
cmd2=(ffmpeg -hide_banner -v error -y -i "$f2" -ss "$t" -frames:v 1 -vf "scale=$scale,format=yuv420p" "$shot2")
vprint "${cmd1[@]}"
"${cmd1[@]}" || { echo "Failed to extract snapshot from $f1 at ${t}s." >&2; exit 3; }
vprint "${cmd2[@]}"
"${cmd2[@]}" || { echo "Failed to extract snapshot from $f2 at ${t}s." >&2; exit 3; }

# Sanity check: files exist and non-empty
if [[ ! -s "$shot1" || ! -s "$shot2" ]]; then
  echo "Snapshot extraction produced empty files. Try a different time or scale." >&2
  [[ $verbose -eq 1 ]] && ls -l "$tmpdir" >&2
  exit 3
fi

# 2) Compare with SSIM (+ PSNR). Different ffmpeg builds print slightly different lines.
# We send stderr to $log and parse multiple possible formats.
cmp_cmd=(ffmpeg -hide_banner -v info -i "$shot1" -i "$shot2" -lavfi "ssim;[0:v][1:v]psnr" -f null -)
vprint "${cmp_cmd[@]}"
"${cmp_cmd[@]}" > /dev/null 2> "$log" || true

# Try to parse SSIM "All:" first, fallback to Y channel, then average/mean lines.
ssim_all="$(grep -Eo 'All:[0-9]+(\.[0-9]+)?' "$log" | head -n1 | cut -d: -f2 || true)"
[[ -z "$ssim_all" ]] && ssim_all="$(grep -Eo 'SSIM [^ ]* All:[0-9]+(\.[0-9]+)?' "$log" | awk -F'All:' '{print $2}' | head -n1 || true)"
[[ -z "$ssim_all" ]] && ssim_all="$(grep -Eo 'SSIM Y:[0-9]+(\.[0-9]+)?' "$log" | head -n1 | cut -d: -f2 || true)"

# Parse PSNR average if present (may appear as "average:xx.xx" or "avg:xx.xx").
psnr_all="$(grep -Eo 'average:[0-9]+(\.[0-9]+)?' "$log" | head -n1 | cut -d: -f2 || true)"
[[ -z "$psnr_all" ]] && psnr_all="$(grep -Eo 'avg:[0-9]+(\.[0-9]+)?' "$log" | head -n1 | cut -d: -f2 || true)"

# 3) Print summaries of the source files (optional but handy)
summarize() {
  local f="$1"
  local dur vinfo fps w h vcodec rfr
  dur=$(ffprobe -v error -show_entries format=duration -of default=nw=1:nk=1 "$f" 2>/dev/null || echo "")
  vinfo=$(ffprobe -v error -select_streams v:0 -show_entries stream=codec_name,width,height,r_frame_rate -of csv=p=0 "$f" 2>/dev/null | head -n1)
  IFS=',' read -r vcodec w h rfr <<<"$vinfo"
  if [[ "$rfr" == */* ]]; then
    fps=$(awk -v r="$rfr" 'BEGIN{split(r,a,"/"); if(a[2]==0) print 0; else printf "%.3f", a[1]/a[2]}')
  else
    fps=$(awk -v r="$rfr" 'BEGIN{printf "%.3f", r+0}')
  fi
  printf "  Duration: %.3fs | Video: %s, %sx%s @ %s fps\n" "${dur:-0}" "${vcodec:-?}" "${w:-?}" "${h:-?}" "${fps:-?}"
}

echo "Snapshot comparison at t=${t}s (scaled to ${scale}):"
echo "File 1: $f1"; summarize "$f1"
echo "File 2: $f2"; summarize "$f2"
echo

# 4) Handle cases where parsing failed
if [[ -z "$ssim_all" && -z "$psnr_all" ]]; then
  echo "Could not parse SSIM/PSNR from ffmpeg output."
  [[ $verbose -eq 1 ]] && { echo "--- ffmpeg log ---"; cat "$log"; echo "-------------------"; }
  echo "Tips: try a different time (e.g. 5 or 30), or a different scale (e.g. 640:-1)."
  exit 0
fi

# 5) Report metrics and give a verdict
printf "Metrics:\n"
[[ -n "$ssim_all" ]] && printf "  SSIM (All): %s   (1.00 = identical)\n" "$ssim_all"
[[ -n "$psnr_all" ]] && printf "  PSNR avg:   %s dB\n" "$psnr_all"
echo

# Verdict from SSIM (fallback to PSNR if SSIM missing)
verdict="Different"
if [[ -n "$ssim_all" ]]; then
  verdict=$(awk -v s="$ssim_all" 'BEGIN{
    if (s+0 >= 0.97) print "Very close (near-identical)";
    else if (s+0 >= 0.90) print "Close";
    else print "Different";
  }')
elif [[ -n "$psnr_all" ]]; then
  verdict=$(awk -v p="$psnr_all" 'BEGIN{
    if (p+0 >= 40) print "Very close (near-identical)";
    else if (p+0 >= 30) print "Close";
    else print "Different";
  }')
fi
echo "Verdict: $verdict"

# Optional: show where log is if verbose
[[ $verbose -eq 1 ]] && { echo "(Log at: $log)"; }