#!/bin/bash

# omarchy:summary=Measure live disk read and write speed
# omarchy:args=[target-dir]

set -e

if [[ -n ${1:-} && ! -d $1 ]]; then
  echo "Usage: omarchy-disk-speedtest [target-dir]" >&2
  exit 2
fi

target_dir="${1:-${XDG_CACHE_HOME:-$HOME/.cache}/omarchy}"
phase_seconds=8
parallel=4
chunk_mb=4
file_mb=256

mkdir -p "$target_dir"

worker_pids=()
chunk_file=""
test_files=()

stop_workers() {
  local pid
  for pid in "${worker_pids[@]}"; do
    [[ -n $pid ]] || continue
    pkill -TERM -P "$pid" 2>/dev/null || true
    kill "$pid" 2>/dev/null || true
  done
  for pid in "${worker_pids[@]}"; do
    [[ -n $pid ]] || continue
    wait "$pid" 2>/dev/null || true
  done
  worker_pids=()
}

alive_workers() {
  local pid count=0
  for pid in "${worker_pids[@]}"; do
    kill -0 "$pid" 2>/dev/null && count=$((count + 1))
  done
  echo "$count"
}

cleanup() {
  # Unlink before stopping the workers, so even a cleanup cut short by an
  # impatient SIGKILL has already taken the names off the filesystem. A live
  # write worker's next dd pass recreates its file by name, so sweep again
  # once they are gone.
  rm -f ${chunk_file:+"$chunk_file"} "${test_files[@]}"
  stop_workers
  rm -f ${chunk_file:+"$chunk_file"} "${test_files[@]}"
}
# Armed before any scratch file exists, so a failed preflight check below
# cannot leak them.
trap cleanup EXIT
trap 'exit 143' TERM INT

# Exclusive per-invocation scratch files: predictable names could clobber a
# user's file, follow a planted symlink, or let overlapping runs delete each
# other's active files out from under the measurement. Each worker gets its
# own on-disk file so the phases run at a queue depth the device can actually
# stretch out on, like the network test's parallel curl workers.
#
# The files are marked NOCOW where the filesystem supports it (btrfs), which
# turns off copy-on-write, checksums, and compression for them. That is what
# makes O_DIRECT truly direct on btrfs -- with checksums on it silently falls
# back to the page cache -- and it makes every rewrite land in place instead
# of churning the extent allocator, which run-to-run reproducibility depends
# on.
chunk_file=$(mktemp /dev/shm/omarchy-disk-speedtest-XXXXXX.src)
for (( i = 0; i < parallel; i++ )); do
  file=$(mktemp "$target_dir/disk-speedtest-XXXXXX.dat")
  chattr +C "$file" 2>/dev/null || true
  test_files+=("$file")
done

format_rate() {
  awk -v value="$1" 'BEGIN {
    if (value <= 0) print "0.0"
    else if (value < 10) printf "%.1f\n", value
    else printf "%.0f\n", value
  }'
}

# Resolve the block device backing the target directory, so throughput can be
# sampled from its kernel I/O counters the same way the network speed test
# samples the interface counters.
source_dev=$(findmnt -no SOURCE --target "$target_dir" 2>/dev/null)
source_dev=${source_dev%%\[*} # Strip btrfs subvolume suffix: /dev/sda2[/@home]

if [[ $source_dev != /dev/* ]]; then
  echo "Cannot find a disk behind $target_dir" >&2
  exit 1
fi

dev=$(readlink -f "$source_dev")
dev=${dev##*/}

if [[ ! -r /sys/class/block/$dev/stat ]]; then
  echo "No I/O statistics for $dev" >&2
  exit 1
fi

available_mb=$(df --output=avail -m "$target_dir" | tail -1 | tr -d ' ')
if (( available_mb < parallel * file_mb * 2 )); then
  echo "Need at least $((parallel * file_mb * 2))MB free on $target_dir" >&2
  exit 1
fi

# Name the physical disk under test, walking dm-crypt/LVM layers and the
# partition table up to the whole device that carries the hardware model.
disk=$dev
while slave=$(ls "/sys/class/block/$disk/slaves" 2>/dev/null | head -1); [[ -n $slave ]]; do
  disk=$slave
done
if [[ -f /sys/class/block/$disk/partition ]]; then
  parent=$(readlink -f "/sys/class/block/$disk")
  parent=${parent%/*}
  disk=${parent##*/}
fi
model=$(lsblk -dno MODEL "/dev/$disk" 2>/dev/null | sed 's/^ *//; s/ *$//')
echo "disk ${model:-$disk}"

# The stress data must be incompressible so nothing between the write call
# and the flash can shrink it. Staging a urandom chunk in RAM also keeps the
# source out of the measurement -- reading tmpfs is a memcpy.
dd if=/dev/urandom of="$chunk_file" bs=${chunk_mb}M count=$((file_mb / chunk_mb)) status=none

# Workers loop only while the main script lives: if cleanup ever loses the
# race with a kill, an orphaned worker finishes its current pass and stops
# instead of hammering the disk forever.
write_worker() {
  local file=$1
  while kill -0 $$ 2>/dev/null; do
    dd if="$chunk_file" of="$file" bs=${chunk_mb}M oflag=direct conv=notrunc status=none 2>/dev/null || return
  done
}

read_worker() {
  local file=$1
  while kill -0 $$ 2>/dev/null; do
    dd if="$file" of=/dev/null bs=${chunk_mb}M iflag=direct status=none 2>/dev/null || return
  done
}

device_sectors() {
  local -a stats
  read -r -a stats < "/sys/class/block/$dev/stat"
  if [[ $1 == "read" ]]; then
    echo "${stats[2]}"
  else
    echo "${stats[6]}"
  fi
}

run_phase() {
  local phase=$1
  local file before after deadline rate alive samples=0
  local baseline_sectors baseline_time end_time

  for file in "${test_files[@]}"; do
    "${phase}_worker" "$file" 2>/dev/null &
    worker_pids+=("$!")
  done

  before=$(device_sectors "$phase")
  deadline=$((SECONDS + phase_seconds))

  while (( SECONDS < deadline )) && (( $(alive_workers) > 0 )); do
    sleep 1
    after=$(device_sectors "$phase")
    end_time=$EPOCHREALTIME
    rate=$(awk -v before="$before" -v after="$after" 'BEGIN {
      if (after < before) print 0
      else print (after - before) * 512 / 1000000
    }')
    echo "$phase $(format_rate "$rate")"
    samples=$((samples + 1))
    # The first second is warm-up -- governor ramp, crypt workers spinning
    # up -- so the steady-state average starts after it.
    if (( samples == 1 )); then
      baseline_sectors=$after
      baseline_time=$end_time
    fi
    before=$after
  done

  # The workers only stop on their own when dd fails (quota, I/O error, full
  # disk), so any worker gone before the deadline is a failed measurement,
  # not a finished one.
  alive=$(alive_workers)
  stop_workers
  if (( alive < parallel )); then
    echo "Disk $phase test failed before finishing" >&2
    exit 1
  fi

  # The figure the dial settles on is the steady-state mean over the whole
  # phase, not whatever rate the final second happened to catch.
  if (( samples > 1 )); then
    rate=$(awk -v before="$baseline_sectors" -v after="$after" -v start="$baseline_time" -v end="$end_time" 'BEGIN {
      secs = end - start
      if (secs <= 0 || after < before) print 0
      else print (after - before) * 512 / 1000000 / secs
    }')
    echo "$phase $(format_rate "$rate")"
  fi
}

# The read phase runs first, so its data must be staged before any measuring
# starts. Direct I/O leaves nothing in the page cache to serve reads from.
for file in "${test_files[@]}"; do
  dd if="$chunk_file" of="$file" bs=${chunk_mb}M oflag=direct conv=notrunc status=none 2>/dev/null &
  worker_pids+=("$!")
done

stage_failed=0
for pid in "${worker_pids[@]}"; do
  wait "$pid" || stage_failed=1
done
worker_pids=()

if (( stage_failed )) || [[ ! -s ${test_files[0]} ]]; then
  echo "Direct disk I/O is not available on $target_dir" >&2
  exit 1
fi

run_phase read
run_phase write