#!/bin/bash
# SPDX-License-Identifier: BSD-3-Clause
# split_screen.sh - compose a side-by-side (split-screen) mp4 from two PPM
# frame sequences captured through the KVMFR ivshmem path (guest publishers:
# tools/cuda-ab/frame_producer.cu or tools/guest-diag/xshm_pub.c; host reader:
# tools/lg/frame_reader.py, which writes frame-%04d.ppm into a directory).
#
# Usage: split_screen.sh <dirA> <dirB> <out.mp4> [fps] [left_w] [left_h]
#   dirA      left-side  frame-%04d.ppm directory (any starting index)
#   dirB      right-side frame-%04d.ppm directory
#   out.mp4   output file (libx264, yuv420p)
#   fps       output frame rate (default 30; pass the measured capture rate,
#             e.g. 10, for real-time playback of ~10 fps captures)
#   left_w    optional left-pane width; the left pane is fitted into
#             left_w x H with letterbox padding (aspect preserved)
#   left_h    optional pane height H; BOTH panes are scaled to height H
#             (default: native height of dirA's first frame). Widths follow
#             each input's aspect ratio (forced even for yuv420p).
#
# Unequal sequence lengths: the shorter side is padded by repeating its last
# frame. hstack is a framesync filter whose repeatlast=1 default holds the
# shorter input's last frame until the longer input ends; shortest=0 is
# passed explicitly to make that the contract (verified on ffmpeg 6.1.1 for
# both input orders). No symlink/tpad workaround is needed.
#
# Fails loudly (exit 1) on: missing ffmpeg/libx264, missing input dirs, or an
# input dir with no frame-*.ppm files. Requires: ffmpeg (libx264); ffprobe is
# used for the final summary when present.
set -euo pipefail

usage() {
    sed -n '3,28p' "$0" | sed 's/^# \{0,1\}//'
}

die() { echo "split_screen.sh: ERROR: $*" >&2; exit 1; }

# grep -q would SIGPIPE the ffmpeg producer under pipefail; capture first.
ffmpeg_has() {  # $1 = encoder/filter token
    local out
    out=$(ffmpeg -hide_banner "$2" 2>/dev/null) || true
    case "$out" in *" $1 "*) return 0 ;; *) return 1 ;; esac
}

# Read width/height from the PPM P6 header (line 2: "<w> <h>").
ppm_dims() {
    local line w h
    line=$(sed -n '2p' "$1")
    read -r w h _ <<< "$line"
    [ -n "${w:-}" ] && [ -n "${h:-}" ] || return 1
    echo "$w $h"
}

# Scan a frame dir: sets PD_COUNT PD_FIRST PD_LAST PD_W PD_H PD_FILE
probe_dir() {
    local dir="$1" label="$2" files f dim
    [ -d "$dir" ] || die "$label directory not found: $dir"
    mapfile -t files < <(find "$dir" -maxdepth 1 -name 'frame-*.ppm' -type f \
        -printf '%f\n' | sort)
    PD_COUNT=${#files[@]}
    [ "$PD_COUNT" -gt 0 ] || die "$label directory has no frame-*.ppm files: $dir"
    f=${files[0]}
    PD_FIRST=${f#frame-}; PD_FIRST=${PD_FIRST%.ppm}; PD_FIRST=$((10#$PD_FIRST))
    f=${files[$((PD_COUNT - 1))]}
    PD_LAST=${f#frame-}; PD_LAST=${PD_LAST%.ppm}; PD_LAST=$((10#$PD_LAST))
    if [ $((PD_LAST - PD_FIRST + 1)) -ne "$PD_COUNT" ]; then
        echo "split_screen.sh: WARN: $label sequence has gaps ($PD_COUNT files," \
             "indices $PD_FIRST..$PD_LAST); ffmpeg stops at the first gap" >&2
    fi
    PD_FILE=$(printf '%s/frame-%04d.ppm' "$dir" "$PD_FIRST")
    dim=$(ppm_dims "$PD_FILE") || die "$label: cannot read PPM header of $PD_FILE"
    PD_W=${dim% *}; PD_H=${dim#* }
}

[ "$#" -ge 3 ] || { usage; exit 2; }
DIR_A=$1 DIR_B=$2 OUT=$3 FPS=${4:-30} LW=${5:-} LH=${6:-}

command -v ffmpeg >/dev/null 2>&1 || die "ffmpeg not found in PATH"
ffmpeg_has libx264 -encoders || die "ffmpeg has no libx264 encoder"

probe_dir "$DIR_A" left;  A_COUNT=$PD_COUNT A_FIRST=$PD_FIRST A_H=$PD_H
probe_dir "$DIR_B" right; B_COUNT=$PD_COUNT B_FIRST=$PD_FIRST

# Pane height H: left_h if given, else dirA's native height. Keep even.
H=${LH:-$A_H}
if [ $((H % 2)) -ne 0 ]; then
    echo "split_screen.sh: WARN: pane height $H rounded down to even for yuv420p" >&2
    H=$((H / 2 * 2))
fi
[ "$H" -ge 2 ] || die "pane height must be >= 2"

# Left pane: fit into left_w x H (letterbox) when left_w given, else scale
# to height H with aspect preserved. Right pane: always scale to height H.
if [ -n "$LW" ]; then
    if [ $((LW % 2)) -ne 0 ]; then
        echo "split_screen.sh: WARN: left_w $LW rounded down to even for yuv420p" >&2
        LW=$((LW / 2 * 2))
    fi
    [ "$LW" -ge 2 ] || die "left_w must be >= 2"
    LEFT="[0:v]scale=${LW}:${H}:force_original_aspect_ratio=decrease,"\
"pad=${LW}:${H}:(ow-iw)/2:(oh-ih)/2,setsar=1[l]"
else
    LEFT="[0:v]scale=-2:${H},setsar=1[l]"
fi
RIGHT="[1:v]scale=-2:${H},setsar=1[r]"
# shortest=0 + framesync repeatlast=1 (default): shorter side's last frame is
# repeated until the longer side ends (see header comment).
FILTER="${LEFT};${RIGHT};[l][r]hstack=shortest=0"

echo "split_screen.sh: left=$DIR_A ($A_COUNT frames from $A_FIRST)" \
     "right=$DIR_B ($B_COUNT frames from $B_FIRST) pane_h=$H fps=$FPS" >&2

ffmpeg -hide_banner -loglevel warning -y \
    -thread_queue_size 512 -framerate "$FPS" -start_number "$A_FIRST" -i "$DIR_A/frame-%04d.ppm" \
    -thread_queue_size 512 -framerate "$FPS" -start_number "$B_FIRST" -i "$DIR_B/frame-%04d.ppm" \
    -filter_complex "$FILTER" \
    -c:v libx264 -preset medium -crf 20 -pix_fmt yuv420p \
    -r "$FPS" -movflags +faststart \
    "$OUT"

[ -s "$OUT" ] || die "ffmpeg produced no output: $OUT"
if command -v ffprobe >/dev/null 2>&1; then
    ffprobe -v error -count_frames -select_streams v:0 \
        -show_entries stream=width,height,nb_read_frames \
        -of default=nw=1 "$OUT" | sed 's/^/split_screen.sh: /' >&2
fi
echo "SPLIT_DONE $OUT" >&2