#!/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: a watchdog checks ffmpeg's frame counter
# and kills/restarts if video freezes (DRM/VAAPI contention from
# other apps using the GPU, e.g. video calls).

set -uo pipefail

RECEIVER_IP="${1:-mcrndeb}"
PORT="${2:-4444}"
STALL_TIMEOUT=10  # seconds with no frame 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)
trap 'rm -f "$PROGRESS_FILE"' EXIT

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"
    )

    "${args[@]}" &
    echo $!
}

get_frame_count() {
    # -progress file writes key=value pairs; frame= is the video frame counter
    grep -oP '^frame=\K[0-9]+' "$PROGRESS_FILE" 2>/dev/null | tail -1
}

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

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

    last_frame=0
    stall_since=$SECONDS

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

        cur_frame=$(get_frame_count)
        cur_frame=${cur_frame:-0}

        if (( cur_frame > last_frame )); then
            last_frame=$cur_frame
            stall_since=$SECONDS
        fi

        if (( SECONDS - stall_since > STALL_TIMEOUT )); then
            echo "Video stalled at frame $last_frame 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