#!/bin/bash
# CHT Sender: muxed video + audio over TCP/mpegts
# Source: Wayland (kmsgrab) + desktop audio + webcam mic
# Usage: ./stream_av.sh [RECEIVER_IP] [PORT]
#
# Requires: sudo for kmsgrab, PulseAudio for audio capture
# Audio is non-blocking (monitor source = passive tap)
#
# Auto-restarts on stall: watchdog checks output bytes + frame counter.
# Also restarts immediately on fatal kmsgrab errors (DRM plane format
# change from fullscreen / direct-scanout).

set -uo pipefail

RECEIVER_IP="${1:-mcrndeb}"
PORT="${2:-4444}"
STALL_TIMEOUT=10  # seconds with no progress before restart

# Let root access the user's PulseAudio session
REAL_UID="${SUDO_UID:-$(id -u)}"
export PULSE_SERVER="unix:/run/user/${REAL_UID}/pulse/native"

# Find the default sink's monitor source (desktop audio - what you hear)
MONITOR=$(PULSE_SERVER="$PULSE_SERVER" pactl info 2>/dev/null | grep "Default Sink" | awk '{print $3}').monitor
# Webcam mic - find by partial match (serial number varies)
WEBCAM_MIC=$(PULSE_SERVER="$PULSE_SERVER" pactl list short sources 2>/dev/null | grep -i "C922" | awk '{print $2}' || true)

echo "Monitor source: $MONITOR"
echo "Webcam mic: ${WEBCAM_MIC:-not found}"
echo "Streaming to: ${RECEIVER_IP}:${PORT}"

# Raise fd limit for long sessions (DMA-BUF fds from kmsgrab)
ulimit -n 65536

PROGRESS_FILE=$(mktemp)
FFLOG=$(mktemp)
FFPID=""

cleanup() {
    [ -n "$FFPID" ] && kill "$FFPID" 2>/dev/null && wait "$FFPID" 2>/dev/null
    rm -f "$PROGRESS_FILE" "$FFLOG"
}
trap cleanup EXIT INT TERM

start_ffmpeg() {
    local args=(
        ffmpeg
        -init_hw_device drm=drm:/dev/dri/card0
        -init_hw_device vaapi=va@drm
        -thread_queue_size 64 -device /dev/dri/card0 -f kmsgrab -framerate 30 -i -
        -thread_queue_size 1024 -f pulse -i "$MONITOR"
    )

    if [ -n "$WEBCAM_MIC" ]; then
        args+=(-thread_queue_size 1024 -f pulse -i "$WEBCAM_MIC")
        args+=(-filter_complex "[1:a][2:a]amix=inputs=2:duration=longest[aout]")
        args+=(-map 0:v -map "[aout]")
    fi

    args+=(
        -vf 'hwmap=derive_device=vaapi,scale_vaapi=w=1920:h=1080:format=nv12,fps=30'
        -c:v h264_vaapi -qp 20 -g 30 -keyint_min 30 -bf 0
        -c:a aac -b:a 128k
        -max_muxing_queue_size 64
        -flush_packets 1 -fflags nobuffer -muxdelay 0 -muxpreload 0
        -f mpegts "tcp://${RECEIVER_IP}:${PORT}"
        -hide_banner -progress "$PROGRESS_FILE"
    )

    > "$FFLOG"
    "${args[@]}" 2>"$FFLOG" &
    echo $!
}

get_progress_val() {
    grep -oP "^${1}=\K[0-9]+" "$PROGRESS_FILE" 2>/dev/null | tail -1
}

fatal_kmsgrab_error() {
    grep -q "framebuffer format changed\|Error during demuxing" "$FFLOG" 2>/dev/null
}

while true; do
    echo "--- Starting sender $(date) ---"
    > "$PROGRESS_FILE"

    FFPID=$(start_ffmpeg)
    echo "ffmpeg started: pid=$FFPID"

    last_bytes=0
    last_frame=0
    stall_since=$SECONDS

    while kill -0 "$FFPID" 2>/dev/null; do
        sleep 1

        # Immediate restart on fatal kmsgrab errors (DRM plane format change)
        if fatal_kmsgrab_error; then
            echo "Fatal kmsgrab error — restarting immediately"
            kill "$FFPID" 2>/dev/null
            wait "$FFPID" 2>/dev/null
            break
        fi

        cur_bytes=$(get_progress_val total_size)
        cur_bytes=${cur_bytes:-0}
        cur_frame=$(get_progress_val frame)
        cur_frame=${cur_frame:-0}

        # Either metric advancing counts as healthy:
        #   total_size: catches TCP output stalls (muxer blocked)
        #   frame:      catches kmsgrab/encoder stalls (audio keeps total_size ticking)
        if (( cur_bytes > last_bytes || cur_frame > last_frame )); then
            last_bytes=$cur_bytes
            last_frame=$cur_frame
            stall_since=$SECONDS
        fi

        if (( SECONDS - stall_since > STALL_TIMEOUT )); then
            echo "Stream stalled at frame ${last_frame} / ${last_bytes}B for ${STALL_TIMEOUT}s — killing ffmpeg"
            kill "$FFPID" 2>/dev/null
            wait "$FFPID" 2>/dev/null
            break
        fi
    done

    if ! kill -0 "$FFPID" 2>/dev/null; then
        wait "$FFPID" 2>/dev/null
        rc=$?
        if (( rc == 0 )); then
            echo "ffmpeg exited cleanly"
            break
        fi
    fi

    echo "Restarting in 2s..."
    sleep 2
done