#!/bin/bash
# Compares live launch with experimental prewarm generation; the UI does not adopt it.
set -u
# Hard rule 9's guard, which owns FIXTURE_ROOT. This was the last script in the tree writing without
# it, found by a class sweep rather than by a failure.
. "$(dirname "$0")/../tools/flea-sandbox-guard"


export PATH="$HOME/.local/bin:$PATH"
eval "$(omarchy-drive env)"

root=$(cd "$(dirname "$0")/.." && pwd) || exit 1
bin=${BIN:-$root/target/release/flea}
ui=${UI:-$root/ui}
dir=${DIR:-$FIXTURE_ROOT/flea-bench-btrfs}
prewarm_path=${XDG_RUNTIME_DIR:-/run/user/$(id -u)}/flea-first-paint-$$.json
# The one path this tool deletes, checked absolute and two components deep before either rm runs.
case $prewarm_path in
  /*/*) ;;
  *) echo "flea-first-paint: refusing a prewarm path that is not absolute: '$prewarm_path'" >&2; exit 1 ;;
esac
first=${FIRST:-350}
iterations=${ITERATIONS:-5}
log=${LOG:-/tmp/flea-first-paint.log}
expected_total=${TOTAL:-100000}
expected_row=${ROW:-some_file_name_0.txt|file|0 B}
timeout_seconds=${TIMEOUT:-15}
shell_pid=
identity_checks=100
identity_poll_seconds=0.01

fail() {
    echo "flea-first-paint: $*"
    exit 1
}

is_flea_process() {
    local pid=$1
    [ -r "/proc/$pid/cmdline" ] || return 1
    tr '\0' ' ' < "/proc/$pid/cmdline" | grep -Fq -- "$ui"
}

wait_for_flea_process() {
    local pid=$1
    local check
    for ((check = 0; check < identity_checks; check++)); do
        is_flea_process "$pid" && return 0
        kill -0 "$pid" 2>/dev/null || return 1
        sleep "$identity_poll_seconds"
    done
    return 1
}

assert_no_flea() {
    local pid
    for pid in $(pgrep -x qs 2>/dev/null); do
        if is_flea_process "$pid"; then
            fail "Flea is already running as pid $pid"
        fi
    done
}

stop_shell() {
    local pid=$1
    if kill -0 "$pid" 2>/dev/null; then
        is_flea_process "$pid" || fail "pid $pid is not the Flea process"
        kill "$pid" || fail "could not stop pid $pid"
    fi
    wait "$pid" 2>/dev/null || true
    shell_pid=
}

cleanup() {
    if [ -n "$shell_pid" ] && kill -0 "$shell_pid" 2>/dev/null; then
        if wait_for_flea_process "$shell_pid"; then
            kill "$shell_pid" 2>/dev/null || true
        fi
        wait "$shell_pid" 2>/dev/null || true
    fi
    rm -f "$prewarm_path"
}

wait_for_row() {
    local pid=$1
    local deadline=$((SECONDS + timeout_seconds))
    local actual
    while ((SECONDS <= deadline)); do
        actual=$(omarchy-drive ipc -p "$ui/boot" flea rowAt 0 2>/dev/null) || true
        if [ "$actual" = "$expected_row" ]; then
            return 0
        fi
        kill -0 "$pid" 2>/dev/null || return 1
        sleep 0.01
    done
    return 1
}

run_once() {
    local label=$1
    local generate_prewarm=$2
    local start
    local pid
    local elapsed

    rm -f "$prewarm_path"
    start=$(date +%s%3N)
    if [ "$generate_prewarm" = yes ]; then
        "$bin" --prewarm "$dir" "$first" "$prewarm_path" || fail "prewarm failed"
    fi

    # Through flea --gui, not a bare qs: the launcher hands over the first-paint colour and trades the platform theme.
    if [ "$generate_prewarm" = yes ]; then
        QSG_RHI_BACKEND="${QSG_RHI_BACKEND:-vulkan}" FLEA_PREWARM="$prewarm_path" FLEA_UI="$ui" FLEA_BIN="$bin" \
            "$bin" --gui "$dir" >"$log" 2>&1 &
    else
        env -u FLEA_PREWARM QSG_RHI_BACKEND="${QSG_RHI_BACKEND:-vulkan}" FLEA_UI="$ui" FLEA_BIN="$bin" \
            "$bin" --gui "$dir" >"$log" 2>&1 &
    fi
    shell_pid=$!
    pid=$shell_pid
    wait_for_flea_process "$pid" || fail "launched pid $pid is not Flea"
    if ! omarchy-drive wait window flea --timeout "$timeout_seconds" >/dev/null; then
        stop_shell "$pid"
        fail "the shell exited before a window appeared"
    fi
    if ! omarchy-drive windows --json | grep -q '"class":"com.thisisgm.flea"'; then
        stop_shell "$pid"
        fail "the mapped window is not Flea"
    fi
    if ! omarchy-drive wait ipc -p "$ui/boot" flea total "$expected_total" --timeout "$timeout_seconds" >/dev/null; then
        stop_shell "$pid"
        fail "the expected total never reached the window"
    fi
    if ! wait_for_row "$pid"; then
        stop_shell "$pid"
        fail "the expected first row never reached the window"
    fi
    elapsed=$(( $(date +%s%3N) - start ))
    printf '  %-22s %4d ms\n' "$label" "$elapsed"
    stop_shell "$pid"
}

[ -x "$bin" ] || fail "$bin is missing, run cargo build --release"
[ -d "$dir" ] || fail "$dir is missing"
[ "$(findmnt -T "$dir" -no FSTYPE)" != tmpfs ] || fail "$dir is on tmpfs"
case "$iterations" in
    ''|*[!0-9]*|0) fail "ITERATIONS must be a positive integer" ;;
esac

assert_no_flea
trap cleanup EXIT

echo "=== first paint, $dir ==="
echo "note: the experimental arm generates a prewarm file that production ignores"
for ((iteration = 1; iteration <= iterations; iteration++)); do
    echo "iteration $iteration"
    run_once "live listing" no
    run_once "prewarm generated" yes
done
