#!/usr/bin/env bash
# Probe the registry's replicate classes against real decoders.
#
# `codecs::capability_replicate_class` says, per capability, whether every
# conforming decoder outputs the identical reconstruction (`standard-defined`,
# resting on a conformance clause that requires it) or whether conforming
# decoders may differ (`standard-tolerance`: the standard states a criterion
# around a reference; `standard-unclassified`: its conformance was not read).
# This script is a falsification check on that claim where this host has more
# than one independent decoder for a format — it does not establish any class:
# it encodes a short synthetic source, decodes it with each decoder to raw
# samples in one common layout, and compares the outputs byte for byte.
#
# Verdicts, per row:
#   identical      every decoder gave the same bytes.
#   differ         decoders disagree; the count and the largest per-sample
#                  difference are printed.
#   skipped        fewer than two decoders (or no encoder) on this host. A
#                  skipped row is not a pass.
#
# The script FAILS (exit 1) only when a `standard-defined` capability's
# decoders disagree — that would contradict the registry. A `standard-tolerance`
# or `standard-unclassified` capability may come out identical on a given
# stream; the class permits that and claims no difference, so it is reported,
# not failed.
#
# Decoders used when present: ffmpeg's own decoders; libdav1d, libvpx and
# libopus through ffmpeg; macOS VideoToolbox hardware decode (forced onto
# hardware surfaces, so it fails rather than falling back to software);
# macOS AudioToolbox AAC; libjpeg-turbo `djpeg`; OpenJPEG `opj_decompress`;
# libFLAC `flac`; Core Audio `afconvert`.
#
#   ./tools/wai_replicate_probe.sh
#
# The CLI is built with `cargo build --bin wai` (default features, so it needs
# libjxl); CARGO_TARGET_DIR is honoured.

set -uo pipefail
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
TMP="$(mktemp -d)"
trap 'rm -rf "$TMP"' EXIT
have() { command -v "$1" >/dev/null 2>&1; }
have ffmpeg || { echo "ffmpeg not found — nothing to probe"; exit 1; }
have python3 || { echo "python3 not found — needed to compare samples"; exit 1; }

WAI="${CARGO_TARGET_DIR:-$ROOT/wai-rs/target}/debug/wai"
echo "building wai CLI…"
( cd "$ROOT/wai-rs" && cargo build --bin wai >/dev/null 2>&1 ) || { echo "build failed"; exit 1; }
"$WAI" capabilities > "$TMP/caps.txt" || { echo "wai capabilities failed"; exit 1; }
class_of() { awk -v c="$1" '$1 == c { print $5 }' "$TMP/caps.txt"; }

echo "host: $(uname -sm) · $(sysctl -n machdep.cpu.brand_string 2>/dev/null || uname -p) · $(ffmpeg -version | head -1 | cut -d' ' -f1-3)"

ffdec_has() { ffmpeg -hide_banner -decoders 2>/dev/null | awk '{print $2}' | grep -qx "$1"; }
ffenc_has() { ffmpeg -hide_banner -encoders 2>/dev/null | awk '{print $2}' | grep -qx "$1"; }
FF=(ffmpeg -hide_banner -loglevel error -y)

# Compare decoder outputs. $1 sample width in bytes (1 = u8, 2 = s16le);
# remaining args are name=file pairs. Prints "identical" or "differ …";
# exits 3 when a file is missing or sizes disagree.
compare() {
  python3 - "$@" <<'PY'
import sys, array
w = int(sys.argv[1]); pairs = [a.split("=", 1) for a in sys.argv[2:]]
def load(p):
    b = open(p, "rb").read()
    if w == 1: return b
    a = array.array("h"); a.frombytes(b[: len(b) // 2 * 2]); return a
try:
    data = [(n, load(p)) for n, p in pairs]
except OSError as e:
    print(f"missing output ({e.filename})"); sys.exit(3)
ref_n, ref = data[0]
if any(len(d) != len(ref) for _, d in data):
    print("sizes differ: " + ", ".join(f"{n}={len(d)}" for n, d in data)); sys.exit(3)
worst = []
for n, d in data[1:]:
    diff = [abs(x - y) for x, y in zip(ref, d)]
    k = sum(1 for v in diff if v)
    if k: worst.append(f"{ref_n}≠{n}: {k}/{len(ref)} samples, max {max(diff)}")
print("identical" if not worst else "differ — " + "; ".join(worst))
PY
}

FAIL=0
row() { # $1 capability  $2 sample width  $3.. name=file
  local cap="$1" w="$2"; shift 2
  local class; class="$(class_of "$cap")"
  if [ -z "$class" ] || [ "$class" = "?" ]; then
    printf '%-22s %-22s ✗ no replicate class in `wai capabilities`\n' "$cap" "${class:-missing}"; FAIL=1; return
  fi
  local have_n=0 args=() names=""
  for p in "$@"; do
    [ -s "${p#*=}" ] && { have_n=$((have_n + 1)); args+=("$p"); names="$names${p%%=*} "; }
  done
  if [ "$have_n" -lt 2 ]; then
    printf '%-22s %-22s skipped (decoders available: %s)\n' "$cap" "$class" "${names:-none}"; return
  fi
  local out; out="$(compare "$w" "${args[@]}")"; local rc=$?
  local note=""
  if [ $rc -ne 0 ]; then note="  ✗ comparison error"; FAIL=1
  elif [ "${out%% *}" = "differ" ] && [ "$class" = "standard-defined" ]; then
    note="  ✗ CONTRADICTS standard-defined"; FAIL=1
  elif [ "${out%% *}" = "differ" ]; then note="  (consistent with $class)"
  elif [ "$class" != "standard-defined" ]; then note="  (the class permits this)"
  fi
  printf '%-22s %-22s %s [%s]%s\n' "$cap" "$class" "$out" "${names% }" "$note"
}

# Video: software decode, single-threaded, to planar yuv420p.
vsw() { "${FF[@]}" -threads 1 -c:v "$2" -i "$3" -f rawvideo -pix_fmt yuv420p "$TMP/$1.yuv" 2>/dev/null; }
# VideoToolbox: frames must stay on hardware surfaces, else hwdownload fails.
vvt() { "${FF[@]}" -hwaccel videotoolbox -hwaccel_output_format videotoolbox_vld -i "$2" \
          -vf hwdownload,format=nv12,format=yuv420p -f rawvideo "$TMP/$1.yuv" 2>/dev/null; }
SRC=(-f lavfi -i "testsrc2=size=320x240:rate=25:duration=0.8" -pix_fmt yuv420p)

printf '\n%-22s %-22s %s\n' "capability" "registry class" "decoded outputs"
printf '%-22s %-22s %s\n' "----------------------" "----------------------" "---------------"

if ffenc_has libx264; then
  "${FF[@]}" "${SRC[@]}" -c:v libx264 -bf 2 -f h264 "$TMP/v.264" 2>/dev/null
  vsw avc_ff h264 "$TMP/v.264"; vvt avc_vt "$TMP/v.264"
fi
row wai.video.avc 1 "ffmpeg=$TMP/avc_ff.yuv" "videotoolbox=$TMP/avc_vt.yuv"

if ffenc_has libx265; then
  "${FF[@]}" "${SRC[@]}" -c:v libx265 -x265-params log-level=error -f hevc "$TMP/v.265" 2>/dev/null
  vsw hevc_ff hevc "$TMP/v.265"; vvt hevc_vt "$TMP/v.265"
fi
row wai.video.hevc 1 "ffmpeg=$TMP/hevc_ff.yuv" "videotoolbox=$TMP/hevc_vt.yuv"

if ffenc_has libvpx-vp9; then
  "${FF[@]}" "${SRC[@]}" -c:v libvpx-vp9 -f ivf "$TMP/v.ivf" 2>/dev/null
  vsw vp9_ff vp9 "$TMP/v.ivf"
  ffdec_has libvpx-vp9 && vsw vp9_vpx libvpx-vp9 "$TMP/v.ivf"
  vvt vp9_vt "$TMP/v.ivf"
fi
row wai.video.vp9 1 "ffmpeg=$TMP/vp9_ff.yuv" "libvpx=$TMP/vp9_vpx.yuv" "videotoolbox=$TMP/vp9_vt.yuv"

AV1ENC=""; for e in libsvtav1 libaom-av1 librav1e; do ffenc_has "$e" && { AV1ENC="$e"; break; }; done
if [ -n "$AV1ENC" ]; then
  "${FF[@]}" "${SRC[@]}" -c:v "$AV1ENC" -f ivf "$TMP/av1.ivf" >/dev/null 2>&1
  ffdec_has libdav1d && vsw av1_dav1d libdav1d "$TMP/av1.ivf"
  ffdec_has libaom-av1 && vsw av1_aom libaom-av1 "$TMP/av1.ivf"
  vvt av1_vt "$TMP/av1.ivf"
fi
row wai.video.av1 1 "dav1d=$TMP/av1_dav1d.yuv" "libaom=$TMP/av1_aom.yuv" "videotoolbox=$TMP/av1_vt.yuv"

# Audio: decode to mono s16le.
TONE=(-f lavfi -i "sine=frequency=440:duration=1:sample_rate=48000" -ac 1)
adec() { "${FF[@]}" -c:a "$2" -i "$3" -f s16le -ac 1 "$TMP/$1.pcm" 2>/dev/null; }

"${FF[@]}" "${TONE[@]}" -c:a aac -f adts "$TMP/a.aac" 2>/dev/null
for d in aac aac_fixed aac_at; do ffdec_has "$d" && adec "aac_$d" "$d" "$TMP/a.aac"; done
row wai.audio.aac 2 "ffmpeg-aac=$TMP/aac_aac.pcm" "ffmpeg-aac_fixed=$TMP/aac_aac_fixed.pcm" "audiotoolbox=$TMP/aac_aac_at.pcm"

if ffenc_has libopus; then
  "${FF[@]}" "${TONE[@]}" -c:a libopus -b:a 64k "$TMP/o.opus" 2>/dev/null
  ffdec_has opus && adec opus_ff opus "$TMP/o.opus"
  ffdec_has libopus && adec opus_lib libopus "$TMP/o.opus"
fi
row wai.audio.opus 2 "ffmpeg-opus=$TMP/opus_ff.pcm" "libopus=$TMP/opus_lib.pcm"

"${FF[@]}" "${TONE[@]}" -c:a pcm_s16le -f s16le "$TMP/flac_src.pcm" 2>/dev/null
"${FF[@]}" -f s16le -ar 48000 -ac 1 -i "$TMP/flac_src.pcm" -c:a flac "$TMP/f.flac" 2>/dev/null
adec flac_ff flac "$TMP/f.flac"
have flac && flac -s -d --force-raw-format --endian=little --sign=signed -o "$TMP/flac_lib.pcm" "$TMP/f.flac" 2>/dev/null
have afconvert && afconvert -f WAVE -d LEI16 "$TMP/f.flac" "$TMP/flac_at.wav" 2>/dev/null \
  && "${FF[@]}" -i "$TMP/flac_at.wav" -f s16le -ac 1 "$TMP/flac_at.pcm" 2>/dev/null
row wai.audio.flac 2 "source=$TMP/flac_src.pcm" "ffmpeg=$TMP/flac_ff.pcm" "libFLAC=$TMP/flac_lib.pcm" "coreaudio=$TMP/flac_at.pcm"

# Still images: grayscale, so no colour conversion enters the comparison.
"${FF[@]}" -f lavfi -i "testsrc2=size=320x240:rate=1:duration=1" -frames:v 1 -pix_fmt gray -f rawvideo "$TMP/src.gray" 2>/dev/null
"${FF[@]}" -f rawvideo -pix_fmt gray -s 320x240 -i "$TMP/src.gray" -f image2 -c:v pgm "$TMP/src.pgm" 2>/dev/null
pgm2raw() { "${FF[@]}" -i "$1" -pix_fmt gray -f rawvideo "$2" 2>/dev/null; }

if have cjpeg; then
  cjpeg -quality 75 -grayscale "$TMP/src.pgm" > "$TMP/g.jpg" 2>/dev/null
  if have djpeg; then
    djpeg -dct int -pnm "$TMP/g.jpg" > "$TMP/j_int.pgm" && pgm2raw "$TMP/j_int.pgm" "$TMP/j_int.gray"
    djpeg -dct float -pnm "$TMP/g.jpg" > "$TMP/j_flt.pgm" && pgm2raw "$TMP/j_flt.pgm" "$TMP/j_flt.gray"
  fi
  "${FF[@]}" -i "$TMP/g.jpg" -pix_fmt gray -f rawvideo "$TMP/j_ff.gray" 2>/dev/null
fi
row wai.image.jpeg 1 "djpeg-int=$TMP/j_int.gray" "djpeg-float=$TMP/j_flt.gray" "ffmpeg=$TMP/j_ff.gray"

if have opj_compress; then
  opj_compress -i "$TMP/src.pgm" -o "$TMP/i.j2k" -I -r 20 >/dev/null 2>&1   # irreversible 9/7
  opj_compress -i "$TMP/src.pgm" -o "$TMP/r.j2k" >/dev/null 2>&1            # reversible 5/3
  for k in i r; do
    have opj_decompress && opj_decompress -i "$TMP/$k.j2k" -o "$TMP/${k}_opj.pgm" >/dev/null 2>&1 \
      && pgm2raw "$TMP/${k}_opj.pgm" "$TMP/${k}_opj.gray"
    "${FF[@]}" -c:v jpeg2000 -i "$TMP/$k.j2k" -pix_fmt gray -f rawvideo "$TMP/${k}_ff.gray" 2>/dev/null
  done
fi
row wai.image.jp2 1 "openjpeg=$TMP/i_opj.gray" "ffmpeg=$TMP/i_ff.gray"
printf '%-22s %-22s ' "  (reversible 5/3)" ""
[ -s "$TMP/r_opj.gray" ] && [ -s "$TMP/r_ff.gray" ] \
  && echo "$(compare 1 "source=$TMP/src.gray" "openjpeg=$TMP/r_opj.gray" "ffmpeg=$TMP/r_ff.gray") [source openjpeg ffmpeg] (lossless path; informational)" \
  || echo "skipped"

echo
if [ $FAIL -ne 0 ]; then
  echo "FAIL — a standard-defined capability's decoders disagreed, a comparison errored, or a probed capability has no registry class."
  exit 1
fi
echo "OK — no standard-defined capability's decoders disagreed on this host (skipped rows prove nothing)."