#!/bin/sh
# mesh-speedtest — quick throughput probe across the wg1 mesh.
#
# Default: from current host, pull a chunk from each peer over ssh and report
# Mbps + a streaming verdict. Bytes traverse the wg1 tunnel; ssh adds a thin
# layer of double-encryption, so numbers run a few % below raw wg throughput
# but are stable for "can I stream from here?" decisions.
#
# Usage:
#   mesh-speedtest                  # download from each peer (64 MiB)
#   mesh-speedtest -s 256           # use 256 MiB chunks (slower, more stable)
#   mesh-speedtest -u               # upload (push to each peer) instead
#   mesh-speedtest -p plum,black    # restrict to specific peers
#   mesh-speedtest --media          # pull from a real media file on black
#                                   #   (validates disk-read + wg path; answers
#                                   #    "can I stream from big disc right now?")
#
# Peers are the wg1 mesh: apricot, black, plum, quinn-vps. Self is skipped.

set -eu

size_mib=64
direction=down
peer_filter=""
media_mode=0
media_peer=black
media_dir=/bigdisk/_/media/unsorted

while [ $# -gt 0 ]; do
    case $1 in
        -s) size_mib=$2; shift 2 ;;
        -u) direction=up; shift ;;
        -p) peer_filter=$2; shift 2 ;;
        --media) media_mode=1; peer_filter=$media_peer; shift ;;
        --media-peer) media_mode=1; media_peer=$2; peer_filter=$2; shift 2 ;;
        --media-dir) media_dir=$2; shift 2 ;;
        -h|--help) sed -n '2,18p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
        *) echo "unknown arg: $1" >&2; exit 2 ;;
    esac
done

all_peers="apricot black plum quinn-vps"

self=$(hostname -s 2>/dev/null || hostname)
case $self in
    plum|*plum*) self=plum ;;
    apricot*) self=apricot ;;
    black*) self=black ;;
    quinn-vps*|*1984*) self=quinn-vps ;;
esac

verdict() {
    mbps=$1
    awk -v m="$mbps" 'BEGIN{
        if (m >= 40) print "4K HDR ok";
        else if (m >= 25) print "4K ok";
        else if (m >= 12) print "1080p high ok";
        else if (m >= 5)  print "1080p ok";
        else if (m >= 3)  print "720p ok";
        else if (m >= 1.5) print "480p / rsync first for HD";
        else print "rsync only";
    }'
}

# Resolve the largest media file under media_dir on the peer; cache after first hit.
media_file_cache=""
resolve_media_file() {
    peer=$1
    if [ -n "$media_file_cache" ]; then
        echo "$media_file_cache"; return
    fi
    f=$(ssh -o ConnectTimeout=5 -o BatchMode=yes "$peer" \
        "find '$media_dir' -maxdepth 4 -type f \( -name '*.mkv' -o -name '*.mp4' -o -name '*.m4v' \) -size +500M 2>/dev/null | head -1")
    media_file_cache=$f
    echo "$f"
}

# Run one direction against one peer. Echo "<peer> <mbps> <secs>" or "<peer> ERR".
probe() {
    peer=$1
    bytes=$((size_mib * 1024 * 1024))
    start=$(date +%s.%N 2>/dev/null || date +%s)
    if [ "$media_mode" = 1 ]; then
        f=$(resolve_media_file "$peer")
        if [ -z "$f" ]; then
            echo "$peer ERR" ; echo "no media file in $media_dir" > /tmp/mesh-speedtest.$$.$peer.err ; return
        fi
        # Read first N MiB of a real file: tests disk + wg + ssh path together.
        ssh -o ConnectTimeout=5 -o BatchMode=yes "$peer" \
            "dd if='$f' bs=1M count=$size_mib iflag=fullblock status=none" \
            > /dev/null 2>/tmp/mesh-speedtest.$$.$peer.err || { echo "$peer ERR"; return; }
    elif [ "$direction" = down ]; then
        ssh -o ConnectTimeout=5 -o BatchMode=yes "$peer" \
            "dd if=/dev/zero bs=1M count=$size_mib status=none" \
            > /dev/null 2>/tmp/mesh-speedtest.$$.$peer.err || { echo "$peer ERR"; return; }
    else
        dd if=/dev/zero bs=1M count=$size_mib status=none 2>/dev/null \
            | ssh -o ConnectTimeout=5 -o BatchMode=yes "$peer" \
              "dd of=/dev/null bs=1M status=none" \
            2>/tmp/mesh-speedtest.$$.$peer.err || { echo "$peer ERR"; return; }
    fi
    end=$(date +%s.%N 2>/dev/null || date +%s)
    secs=$(awk -v s="$start" -v e="$end" 'BEGIN{d=e-s; if(d<0.001)d=0.001; printf "%.3f",d}')
    mbps=$(awk -v b="$bytes" -v t="$secs" 'BEGIN{printf "%.1f", (b*8)/(t*1000000)}')
    echo "$peer $mbps $secs"
}

# Build peer list.
peers=""
if [ -n "$peer_filter" ]; then
    peers=$(echo "$peer_filter" | tr ',' ' ')
else
    for p in $all_peers; do
        [ "$p" = "$self" ] && continue
        peers="$peers $p"
    done
fi

if [ "$media_mode" = 1 ]; then
    label="media-read from"
    f=$(resolve_media_file "$media_peer")
    [ -n "$f" ] && extra=" file=$(basename "$f")" || extra=" (no eligible file in $media_dir)"
else
    label=$([ "$direction" = down ] && echo "download from" || echo "upload to")
    extra=""
fi
printf "mesh-speedtest: %s peers, %d MiB chunks (self=%s)%s\n\n" "$label" "$size_mib" "$self" "$extra"
printf "%-12s %10s %8s   %s\n" "peer" "Mbps" "secs" "verdict"
printf "%-12s %10s %8s   %s\n" "----" "----" "----" "-------"

# Run probes sequentially; parallel would skew Mbps via shared uplink.
for p in $peers; do
    out=$(probe "$p")
    case $out in
        *ERR)
            err=$(cat /tmp/mesh-speedtest.$$.$p.err 2>/dev/null | head -1)
            printf "%-12s %10s %8s   %s\n" "$p" "-" "-" "ssh failed: ${err:-unreachable}"
            ;;
        *)
            peer_n=$(echo "$out" | awk '{print $1}')
            mbps=$(echo "$out" | awk '{print $2}')
            secs=$(echo "$out" | awk '{print $3}')
            v=$(verdict "$mbps")
            printf "%-12s %10s %8s   %s\n" "$peer_n" "$mbps" "$secs" "$v"
            ;;
    esac
    rm -f /tmp/mesh-speedtest.$$.$p.err
done

# Average across reachable peers (informational only — wildly different paths).
echo ""