#!/usr/bin/env bash
# Regression test for scpcap's growing-file tailer, run against a real
# remote host over ssh. Exercises two distinct writer behaviors, in both
# transfer directions, so the tailer is tested wherever it actually runs
# (locally, or in the remote --server process):
#
#   1. "normal growth": a well-behaved writer that only ever appends bytes
#      (file length always equals real content) -- the original behavior
#      this tool was built for. Must keep working.
#
#   2. "preallocated growth": a writer that reserves space ahead of its
#      real write cursor (ftruncate-extend, reading back as zero) and fills
#      it in place, only shrinking the file back down to the real size at
#      rotation, immediately before renaming to final. This is the pattern
#      that caused a hang + corrupted destination before the confirmed_len
#      fix in src/chunker.rs / src/tail.rs -- see git log for the writeup.
#
# Requires: a real pcap file to source test data from (default: the first
# *.pcap file found in ~/pcap), and ssh access to a test host with an
# scpcap binary deployable to ~/bin/scpcap (this script builds and deploys
# it before testing).
#
# Usage: scripts/regression_test.sh [pcap-file] [user@host]

set -uo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
SRC_PCAP="${1:-$(ls "$HOME"/pcap/*.pcap 2>/dev/null | head -1)}"
HOST="${2:-eric@locl.sh}"
REMOTE_EXE="bin/scpcap"
BIN="$REPO_ROOT/target/release/scpcap"
LOCAL_TEST="$HOME/scpcap_test"
REMOTE_TEST="scpcap_test"
PAD_BYTES=65536
STAGE_DELAY=0.3
WORK="$(mktemp -d)"
trap 'rm -rf "$WORK"' EXIT

PASS=0
FAIL=0
log()  { echo "[*] $*" >&2; }
pass() { echo "    PASS: $*"; PASS=$((PASS + 1)); }
fail() { echo "    FAIL: $*"; FAIL=$((FAIL + 1)); }

if [ -z "$SRC_PCAP" ] || [ ! -f "$SRC_PCAP" ]; then
    echo "no source pcap file found (looked in ~/pcap/*.pcap); pass one explicitly" >&2
    exit 2
fi
log "source capture: $SRC_PCAP ($(wc -c < "$SRC_PCAP") bytes)"

log "building release binary"
(cd "$REPO_ROOT" && cargo build --release --quiet) || exit 1
log "deploying to $HOST:~/$REMOTE_EXE"
scp -q "$BIN" "$HOST:$REMOTE_EXE" || exit 1

mkdir -p "$LOCAL_TEST/dest"
ssh "$HOST" "mkdir -p $REMOTE_TEST/dest"

# --- Build a small, record-aligned test capture -----------------------------

TRUNCATED_SRC="$WORK/src.pcap"
head -c 8000000 "$SRC_PCAP" > "$WORK/src.raw" 2>/dev/null || cp "$SRC_PCAP" "$WORK/src.raw"
PREFIX_LEN=$(python3 - "$WORK/src.raw" <<'PYEOF'
import struct, sys
with open(sys.argv[1], "rb") as f:
    data = f.read()
magic = struct.unpack("<I", data[0:4])[0]
endian = "<" if magic in (0xA1B2C3D4, 0xA1B23C4D) else ">"
pos = 24
while pos + 16 <= len(data):
    incl_len = struct.unpack(endian + "I", data[pos + 8: pos + 12])[0]
    total = 16 + incl_len
    if pos + total > len(data):
        break
    pos += total
print(pos)
PYEOF
)
dd if="$WORK/src.raw" of="$TRUNCATED_SRC" bs=1 count="$PREFIX_LEN" status=none
log "test capture: $TRUNCATED_SRC ($(wc -c < "$TRUNCATED_SRC") bytes, record-aligned)"

N_SEGMENTS=4
mapfile -t OFFSETS < <(python3 "$REPO_ROOT/scripts/pcap_offsets.py" "$TRUNCATED_SRC" "$N_SEGMENTS")
FINAL_SIZE="${OFFSETS[-1]}"
log "growth stages (cumulative bytes): ${OFFSETS[*]}"

sha() { sha256sum "$1" 2>/dev/null | awk '{print $1}'; }
EXPECTED_SHA="$(sha "$TRUNCATED_SRC")"

# --- Target primitives: "local:<path>" or "remote:<path>" -------------------

t_write() { # target prev len
    local target="$1" prev="$2" len="$3"
    case "$target" in
        local:*)
            local path="${target#local:}"
            python3 -c "
with open('$TRUNCATED_SRC','rb') as s, open('$path','r+b') as d:
    s.seek($prev); d.seek($prev); d.write(s.read($len))
"
            ;;
        remote:*)
            local path="${target#remote:}"
            dd if="$TRUNCATED_SRC" bs=1 skip="$prev" count="$len" status=none | \
                ssh "$HOST" "python3 -c \"
import sys
with open('$path','r+b') as f:
    f.seek($prev)
    f.write(sys.stdin.buffer.read())
\""
            ;;
    esac
}

t_set_len() { # target len
    local target="$1" len="$2"
    case "$target" in
        local:*) python3 -c "open('${target#local:}','r+b').truncate($len)" ;;
        remote:*) ssh "$HOST" "python3 -c \"open('${target#remote:}','r+b').truncate($len)\"" ;;
    esac
}

t_create_empty() { # target
    local target="$1"
    case "$target" in
        local:*) : > "${target#local:}" ;;
        remote:*) ssh "$HOST" ": > ${target#remote:}" ;;
    esac
}

t_rename() { # target_partial target_final
    local partial="$1" final="$2"
    case "$partial" in
        local:*) mv "${partial#local:}" "${final#local:}" ;;
        remote:*) ssh "$HOST" "mv ${partial#remote:} ${final#remote:}" ;;
    esac
}

# --- Test runner --------------------------------------------------------------

# $1 = name, $2 = "normal" | "preallocated", $3 = "local-to-remote" | "remote-to-local"
run_case() {
    local name="$1" mode="$2" direction="$3"
    log "=== $name ($direction, $mode) ==="

    if [ "$direction" = "local-to-remote" ]; then
        local partial="local:$LOCAL_TEST/$name.pcap.partial"
        local final="local:$LOCAL_TEST/$name.pcap"
        local dest="$REMOTE_TEST/$name.pcap"
        rm -f "$LOCAL_TEST/$name.pcap.partial" "$LOCAL_TEST/$name.pcap"
        ssh "$HOST" "rm -f $dest $dest.partial"

        t_create_empty "$partial"
        t_write "$partial" 0 "${OFFSETS[0]}"
        # For preallocated growth, reserve the padded size right after the
        # first real write, before the client ever gets a chance to look.
        if [ "$mode" = preallocated ]; then
            t_set_len "$partial" $((FINAL_SIZE + PAD_BYTES))
        fi

        "$BIN" --remote-exe "$REMOTE_EXE" "$LOCAL_TEST/$name.pcap.partial" "$HOST:$REMOTE_TEST/" &
        local client_pid=$!

        local prev="${OFFSETS[0]}"
        for off in "${OFFSETS[@]:1}"; do
            sleep "$STAGE_DELAY"
            t_write "$partial" "$prev" $((off - prev))
            prev="$off"
        done
        if [ "$mode" = preallocated ]; then
            sleep "$STAGE_DELAY"
            t_set_len "$partial" "$FINAL_SIZE"
        fi
        t_rename "$partial" "$final"

        wait "$client_pid"
        local status=$?
        local actual_sha
        actual_sha="$(ssh "$HOST" "sha256sum $dest 2>/dev/null | awk '{print \$1}'" || true)"
    else
        local partial="remote:$REMOTE_TEST/$name.pcap.partial"
        local final="remote:$REMOTE_TEST/$name.pcap"
        local dest="$LOCAL_TEST/dest/$name.pcap"
        ssh "$HOST" "rm -f $REMOTE_TEST/$name.pcap $REMOTE_TEST/$name.pcap.partial"
        rm -f "$dest" "$dest.partial"

        t_create_empty "$partial"
        t_write "$partial" 0 "${OFFSETS[0]}"
        if [ "$mode" = preallocated ]; then
            t_set_len "$partial" $((FINAL_SIZE + PAD_BYTES))
        fi

        "$BIN" --remote-exe "$REMOTE_EXE" "$HOST:$REMOTE_TEST/$name.pcap.partial" "$LOCAL_TEST/dest/" &
        local client_pid=$!

        local prev="${OFFSETS[0]}"
        for off in "${OFFSETS[@]:1}"; do
            sleep "$STAGE_DELAY"
            t_write "$partial" "$prev" $((off - prev))
            prev="$off"
        done
        if [ "$mode" = preallocated ]; then
            sleep "$STAGE_DELAY"
            t_set_len "$partial" "$FINAL_SIZE"
        fi
        t_rename "$partial" "$final"

        wait "$client_pid"
        local status=$?
        local actual_sha
        actual_sha="$(sha "$dest" || true)"
    fi

    if [ "$status" -eq 0 ] && [ "$actual_sha" = "$EXPECTED_SHA" ]; then
        pass "$name ($direction): exit 0, sha matches"
    else
        fail "$name ($direction): exit=$status expected_sha=$EXPECTED_SHA actual_sha=$actual_sha"
    fi
}

run_case "normal_growth" normal "local-to-remote"
run_case "preallocated" preallocated "local-to-remote"
run_case "normal_growth_remote_src" normal "remote-to-local"
run_case "preallocated_remote_src" preallocated "remote-to-local"

echo
echo "=== summary: $PASS passed, $FAIL failed ==="
exit $((FAIL > 0))