#!/usr/bin/env bash
set -euo pipefail

# wcs-frame-match.sh
# Extract a frame from MP4 at a given timestamp and compare it to a reference image.
# If it matches within a threshold, prefix the MP4 filename (e.g., WCS_...).
# Dependencies: ffmpeg, ImageMagick (IM6/IM7). NOTE: SSIM typically requires ImageMagick 7 (magick); IM6 often lacks SSIM.

# Defaults
TIMESTAMP="1.0"
TIMESTAMP_LIST=""
METRIC="ssim"        # ssim | rmse | psnr
THRESHOLD=""
PREFIX="WCS_"
DRYRUN="false"
KEEP_TEMP="false"
RECURSIVE="false"
ALIGN="scale"        # scale | pad | none
DEBUG="false"
AUTO_SKIP_BLACK="false"
BLACK_PIC_TH="0.98"
BLACK_MIN_D="0.05"   # seconds
ROI_GEOM=""          # e.g., 800x200+50+100

usage() {
  cat <<'EOF'
Usage: wcs-frame-match.sh -s REF.png [options] [files.mp4...]
  -s  Path to reference image (PNG/JPG etc.) [required]
  -t  Timestamp in seconds (float OK), default: 1.0
      You can also provide a comma-separated list via --timestamps or --timestamp, e.g. 0.0,0.5,1.0
  -m  Metric: ssim | rmse | psnr (default: ssim)  [NOTE: SSIM usually needs ImageMagick 7]
  -T  Threshold/tolerance:
        * ssim: 0..1 (higher is better). e.g., 0.95
        * rmse: 0..1 (lower is better).  e.g., 0.01
        * psnr: dB (higher is better).   e.g., 35
      If not set: ssim=0.95, rmse=0.01, psnr=35
  -p  Prefix to add on match (default: WCS_)
  -r  Recursive: process all *.mp4 under current directory
  -a  Frame sizing vs reference: scale | pad | none (default: scale)
  -n  Dry run (no actual rename)
  -k  Keep temporary files
  -d  Debug (print raw ImageMagick output and details)
  -B  Auto-skip initial black segment using ffmpeg blackdetect (default: off)
      --black-pic-th FLOAT   (default: 0.98)
      --black-min-d  SEC     (default: 0.05)
  -R  ROI crop geometry on aligned images, e.g., 800x200+50+100 (focus only this region)
  -h  Help

Examples:
  wcs-frame-match.sh -s ref.png video1.mp4
  wcs-frame-match.sh -s ref.jpg -t 0.8 -m ssim -T 0.93 film.mp4
  wcs-frame-match.sh -s ref.png -m rmse -T 0.005 -r -n
  wcs-frame-match.sh -s ref.png -m psnr -T 35 -a pad *.mp4
  wcs-frame-match.sh -s ref.png -B --timestamps 0,0.5,1 -R 800x200+0+700 -n *.mp4
EOF
}

# ---- Arg parsing
REF_IMAGE=""
USER_SET_THRESHOLD="false"
# Support long flags via getopts "-:" trick
while getopts ":s:t:m:T:p:rkna:hdBR:R:-:" opt; do
  case "$opt" in
    s) REF_IMAGE="$OPTARG" ;;
    t) TIMESTAMP="$OPTARG" ;;
    m) METRIC="$(echo "$OPTARG" | tr '[:upper:]' '[:lower:]')" ;;
    T) THRESHOLD="$OPTARG"; USER_SET_THRESHOLD="true" ;;
    p) PREFIX="$OPTARG" ;;
    r) RECURSIVE="true" ;;
    n) DRYRUN="true" ;;
    k) KEEP_TEMP="true" ;;
    a) ALIGN="$OPTARG" ;;
    d) DEBUG="true" ;;
    B) AUTO_SKIP_BLACK="true" ;;
    R) ROI_GEOM="$OPTARG" ;;
    -)
      case "$OPTARG" in
        debug) DEBUG="true" ;;
        timestamps=*) TIMESTAMP_LIST="${OPTARG#timestamps=}" ;;
        timestamps)
          # support space-separated: --timestamps 0,0.5,1
          val="${!OPTIND}"; OPTIND=$((OPTIND+1)); TIMESTAMP_LIST="$val" ;;
        timestamp=*) TIMESTAMP_LIST="${OPTARG#timestamp=}" ;;
        timestamp)
          # support space-separated: --timestamp 0,0.5,1 OR single value
          val="${!OPTIND}"; OPTIND=$((OPTIND+1)); TIMESTAMP_LIST="$val" ;;
        black-pic-th=*) BLACK_PIC_TH="${OPTARG#black-pic-th=}" ;;
        black-pic-th)
          val="${!OPTIND}"; OPTIND=$((OPTIND+1)); BLACK_PIC_TH="$val" ;;
        black-min-d=*)  BLACK_MIN_D="${OPTARG#black-min-d=}" ;;
        black-min-d)
          val="${!OPTIND}"; OPTIND=$((OPTIND+1)); BLACK_MIN_D="$val" ;;
        *) echo "Unknown long flag --$OPTARG" >&2; usage; exit 1 ;;
      esac
      ;;
    h) usage; exit 0 ;;
    \?) echo "Invalid flag: -$OPTARG" >&2; usage; exit 1 ;;
    :)  echo "Flag -$OPTARG requires a value." >&2; usage; exit 1 ;;
  esac
done
shift $((OPTIND-1))

# ---- Basic checks
if [[ -z "$REF_IMAGE" ]]; then
  echo "Error: reference image must be provided with -s." >&2
  usage; exit 1
fi
if [[ ! -f "$REF_IMAGE" ]]; then
  echo "Error: reference image not found: $REF_IMAGE" >&2
  exit 1
fi

case "$METRIC" in
  ssim|rmse|psnr) ;;
  *) echo "Error: metric must be one of: ssim | rmse | psnr" >&2; exit 1 ;;
esac

# Set default threshold only if user didn't provide one
if [[ "$USER_SET_THRESHOLD" != "true" ]]; then
  case "$METRIC" in
    ssim) THRESHOLD="0.95" ;;
    rmse) THRESHOLD="0.01" ;;
    psnr) THRESHOLD="35" ;;
  esac
fi

# Locate ImageMagick cmds (prefer IM7 'magick compare' for SSIM support)
if command -v magick >/dev/null 2>&1; then
  COMPARE_CMD=(magick compare)
  IDENTIFY=(magick identify)
  CONVERT=(magick convert)
elif command -v compare >/dev/null 2>&1 && command -v identify >/dev/null 2>&1 && command -v convert >/dev/null 2>&1; then
  COMPARE_CMD=(compare)
  IDENTIFY=(identify)
  CONVERT=(convert)
else
  echo "Error: Could not find ImageMagick (compare/identify/convert or magick) in PATH." >&2
  exit 1
fi

# Create temp dir early (needed by ensure_metric_supported)
TMPDIR="$(mktemp -d -t wcsframe-XXXXXXXX)"
cleanup() {
  if [[ "$KEEP_TEMP" != "true" ]]; then
    rm -rf "$TMPDIR"
  else
    echo "Temporary files kept in: $TMPDIR"
  fi
}
trap cleanup EXIT

# Verify metric support for current ImageMagick
ensure_metric_supported() {
  local metric_upper t1 t2 raw status
  metric_upper="$(echo "$METRIC" | tr '[:lower:]' '[:upper:]')"
  t1="$(mktemp -p "$TMPDIR" imtestA_XXXX.png)"
  t2="$(mktemp -p "$TMPDIR" imtestB_XXXX.png)"
  ${CONVERT[@]} -size 2x2 xc:black "$t1"
  ${CONVERT[@]} -size 2x2 xc:white "$t2"
  set +e
  raw=$(${COMPARE_CMD[@]} -metric "$metric_upper" "$t1" "$t2" null: 2>&1)
  status=$?
  set -e
  if echo "$raw" | grep -qi "unrecognized metric type"; then
    echo "[WARNING] Your ImageMagick variant does not support metric '$metric_upper'." >&2
    if [[ "$METRIC" == "ssim" ]]; then
      echo "[WARNING] Falling back to -m rmse (use -m psnr if you prefer)." >&2
      METRIC="rmse"
      # update default threshold for new metric if user didn't set one
      if [[ "$USER_SET_THRESHOLD" != "true" ]]; then THRESHOLD="0.01"; fi
    else
      echo "[ERROR] Choose a supported metric: -m rmse or -m psnr, or install ImageMagick 7 for SSIM." >&2
    fi
  fi
}

ensure_metric_supported

# Read reference dimensions
REF_W="$(${IDENTIFY[@]} -format "%w" "$REF_IMAGE")"
REF_H="$(${IDENTIFY[@]} -format "%h" "$REF_IMAGE")"

# Precompute optional ROI on reference (after alignment size)
REF_ALIGNED="$TMPDIR/ref_aligned.png"
${CONVERT[@]} "$REF_IMAGE" -resize "${REF_W}x${REF_H}!" -colorspace sRGB "$REF_ALIGNED"
if [[ -n "$ROI_GEOM" ]]; then
  REF_CROP="$TMPDIR/ref_crop.png"
  ${CONVERT[@]} "$REF_ALIGNED" -crop "$ROI_GEOM" +repage "$REF_CROP"
else
  REF_CROP="$REF_ALIGNED"
fi

# Collect files
FILES=()
if [[ "$RECURSIVE" == "true" ]]; then
  while IFS= read -r -d '' f; do FILES+=("$f"); done < <(find . -type f -iname "*.mp4" -print0)
fi
if [[ $# -gt 0 ]]; then
  for f in "$@"; do FILES+=("$f"); done
fi

if [[ ${#FILES[@]} -eq 0 ]]; then
  echo "No MP4 files to process. Provide files or use -r." >&2
  exit 1
fi

# Strict metric value parsing
parse_metric_value() {
  local metric="$1"; shift
  local text="$*"
  local val=""
  case "$metric" in
    ssim|rmse)
      # Only accept values in [0..1] (or exactly 1.0)
      val="$(echo "$text" | grep -Eo '\b(0(\.[0-9]+)?|1(\.0+)?)\b' | head -n1)"
      ;;
    psnr)
      # First float in 0..100 (reasonable range)
      val="$(echo "$text" | grep -Eo '[0-9]+(\.[0-9]+)?' | awk '$1>=0 && $1<=100 {print; exit}')"
      ;;
  esac
  echo "$val"
}

# Resize/alignment (with colorspace normalization)
resize_or_align() {
  local in="$1"
  local out="$2"
  case "$ALIGN" in
    scale)
      ${CONVERT[@]} "$in" -resize "${REF_W}x${REF_H}!" -colorspace sRGB "$out"
      ;;
    pad)
      ${CONVERT[@]} "$in" -resize "${REF_W}x${REF_H}" -background black -gravity center -extent "${REF_W}x${REF_H}" -colorspace sRGB "$out"
      ;;
    none)
      ${CONVERT[@]} "$in" -colorspace sRGB "$out"
      ;;
    *) echo "Invalid ALIGN: $ALIGN" >&2; exit 1 ;;
  esac
}

apply_roi_if_any() {
  local in="$1"; local out="$2"
  if [[ -n "$ROI_GEOM" ]]; then
    ${CONVERT[@]} "$in" -crop "$ROI_GEOM" +repage "$out"
  else
    cp -f "$in" "$out"
  fi
}

# Detect initial black segment end using ffmpeg blackdetect
# Returns end time in seconds (float) or 0 if none found
black_end_time() {
  local mp4="$1"
  local out
  out=$(ffmpeg -hide_banner -nostats -t 5 -i "$mp4" -vf "blackdetect=d=${BLACK_MIN_D}:pic_th=${BLACK_PIC_TH}" -an -f null - 2>&1 || true)
  echo "$out" | awk -F'black_end:' '/black_end/ {gsub(/ .*/, "", $2); print $2; exit}' | awk '{print ($1+0)}'
}

# sample a given timestamp and compare; echoes "score matchflag"; returns 0 always
compare_at_time() {
  local mp4="$1"; local ts="$2"
  local frame_png="$TMPDIR/frame.png"
  local frame_aligned="$TMPDIR/frame_aligned.png"
  local frame_roi="$TMPDIR/frame_roi.png"

  ffmpeg -loglevel error -ss "$ts" -i "$mp4" -frames:v 1 -q:v 2 -y "$frame_png"
  resize_or_align "$frame_png" "$frame_aligned"
  apply_roi_if_any "$frame_aligned" "$frame_roi"

  local metric_upper raw status val match="false"
  metric_upper="$(echo "$METRIC" | tr '[:lower:]' '[:upper:]')"
  set +e
  raw=$(${COMPARE_CMD[@]} -metric "$metric_upper" "$REF_CROP" "$frame_roi" null: 2>&1)
  status=$?
  set -e
  [[ "$DEBUG" == "true" ]] && echo "[DEBUG] t=$ts status=$status raw='$raw'" >&2

  val="$(parse_metric_value "$METRIC" "$raw")"
  if [[ -z "$val" ]]; then
    echo "nan false"; return
  fi
  case "$METRIC" in
    ssim|psnr)
      awk -v v="$val" -v th="$THRESHOLD" 'BEGIN{exit !(v >= th)}' && match="true" || true ;;
    rmse)
      awk -v v="$val" -v th="$THRESHOLD" 'BEGIN{exit !(v <= th)}' && match="true" || true ;;
  esac
  echo "$val $match"
}

process_file() {
  local mp4="$1"
  if [[ ! -f "$mp4" ]]; then
    echo "Skipping (not found): $mp4"; return
  fi
  local base="$(basename "$mp4")"
  if [[ "$base" == "$PREFIX"* ]]; then
    echo "Already prefixed, skipping: $mp4"; return
  fi

  # Determine timestamps to test
  local ts_list=()
  if [[ -n "$TIMESTAMP_LIST" ]]; then
    IFS=',' read -r -a ts_list <<< "$TIMESTAMP_LIST"
  else
    ts_list=("$TIMESTAMP")
  fi

  # Optional: skip initial black
  local bend=0
  if [[ "$AUTO_SKIP_BLACK" == "true" ]]; then
    bend="$(black_end_time "$mp4")"
    bend=${bend:-0}
    [[ -z "$bend" ]] && bend=0
    [[ "$DEBUG" == "true" ]] && echo "[DEBUG] black_end=$bend for $mp4"
  fi

  local best_score=""
  local best_ts=""
  local matched="false"

  for ts in "${ts_list[@]}"; do
    # adjust timestamp if it falls within initial black region
    local eff_ts="$ts"
    if [[ "$AUTO_SKIP_BLACK" == "true" ]]; then
      awk -v a="$eff_ts" -v b="$bend" 'BEGIN{exit !(a < b)}' && eff_ts=$(awk -v b="$bend" 'BEGIN{printf "%.3f", b+0.010}') || true
    fi
    read -r score ismatch < <(compare_at_time "$mp4" "$eff_ts")
    [[ "$score" == "nan" ]] && continue
    if [[ "$matched" == "false" && "$ismatch" == "true" ]]; then
      matched="true"; best_score="$score"; best_ts="$eff_ts"; break
    fi
    # track best score even if not matched
    if [[ -z "$best_score" ]]; then best_score="$score"; best_ts="$eff_ts"; else
      case "$METRIC" in
        ssim|psnr) awk -v x="$score" -v y="$best_score" 'BEGIN{exit !(x>y)}' && { best_score="$score"; best_ts="$eff_ts"; } || true ;;
        rmse)      awk -v x="$score" -v y="$best_score" 'BEGIN{exit !(x<y)}' && { best_score="$score"; best_ts="$eff_ts"; } || true ;;
      esac
    fi
  done

  local verdict
  if [[ "$matched" == "true" ]]; then
    verdict="MATCH"
  else
    verdict="no-match"
  fi
  printf "%s  =>  best %s at t=%s (threshold %s)  =>  %s
" \
    "$mp4" "$METRIC=$best_score" "$best_ts" "$THRESHOLD" "$verdict"

  if [[ "$matched" == "true" ]]; then
    local dir newpath
    dir="$(dirname "$mp4")"; newpath="${dir}/${PREFIX}${base}"
    if [[ "$DRYRUN" == "true" ]]; then
      echo "[DRYRUN] mv -- \"$mp4\" \"$newpath\""
    else
      mv -- "$mp4" "$newpath"
    fi
  fi
}

# Collect either recursive or provided files
if [[ ${#FILES[@]} -eq 0 ]]; then
  : # already handled earlier
fi

# Run
for f in "${FILES[@]}"; do
  process_file "$f"
done