#!/usr/bin/env bash
set -euo pipefail
LC_ALL=C

ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
RUN_ID=$(date +"%Y%m%d_%H%M%S")
OUT_DIR="$ROOT_DIR/experiment_runs/$RUN_ID"
mkdir -p "$OUT_DIR"

LOG_FILE="$OUT_DIR/activity.log"

log() {
  echo "[$(date -Iseconds)] $*" | tee -a "$LOG_FILE"
}

count_processes() {
  ls -d /proc/[0-9]* 2>/dev/null | wc -l
}

count_threads_total() {
  awk '/^Threads:/ {sum+=$2} END {print sum+0}' /proc/[0-9]*/status 2>/dev/null || echo 0
}

get_mem_kb() {
  awk '/^MemTotal:/ {t=$2} /^MemAvailable:/ {a=$2} END {print t" "a}' /proc/meminfo
}

sample_process_csv() {
  local pid="$1"
  local csv="$2"
  local start_uptime="$3"
  local max_seconds="$4"
  local proc_step="${5:-0}"

  local have_pidstat=0
  if command -v pidstat >/dev/null 2>&1; then
    have_pidstat=1
  else
    log "pidstat not found; install sysstat for accurate CPU sampling"
  fi

  local prev_cpu_total=0
  local prev_cpu_idle=0
  if [ -r /proc/stat ]; then
    read -r _ c_user c_nice c_system c_idle c_iowait c_irq c_softirq c_steal _ _ < /proc/stat
    prev_cpu_total=$((c_user + c_nice + c_system + c_idle + c_iowait + c_irq + c_softirq + c_steal))
    prev_cpu_idle=$((c_idle + c_iowait))
  fi

  local first_sample=1
  local next_proc_sample=0

  while kill -0 "$pid" 2>/dev/null; do
    local now_uptime elapsed_s
    now_uptime=$(awk '{print $1}' /proc/uptime)
    elapsed_s=$(awk -v s="$start_uptime" -v n="$now_uptime" 'BEGIN{printf "%.3f", (n-s)}')

    if [ "$max_seconds" -gt 0 ]; then
      if awk -v e="$elapsed_s" -v m="$max_seconds" 'BEGIN{exit (e>m)?0:1}'; then
        log "Timeout reached (${max_seconds}s), stopping PID=$pid"
        kill -TERM "-$pid" 2>/dev/null || true
        sleep 2
        kill -KILL "-$pid" 2>/dev/null || true
        break
      fi
    fi

    local proc_count thread_total
    proc_count=$(count_processes || echo 0)
    thread_total=$(count_threads_total || echo 0)

    local should_sample=1
    if [ "$proc_step" -gt 0 ]; then
      if [ "$first_sample" -eq 1 ]; then
        next_proc_sample=$((((proc_count / proc_step) + 1) * proc_step))
        first_sample=0
      fi
      if [ "$proc_count" -lt "$next_proc_sample" ]; then
        should_sample=0
      else
        should_sample=1
        next_proc_sample=$((next_proc_sample + proc_step))
      fi
    fi

    if [ "$should_sample" -eq 1 ]; then
      local load1 load5 load15
      read -r load1 load5 load15 _ < /proc/loadavg

      local mem_total mem_avail mem_used
      read -r mem_total mem_avail < <(get_mem_kb || echo "0 0")
      mem_used=$((mem_total - mem_avail))

    # Use pidstat for CPU and ps for memory/thread metrics.
    local cpu rss vsz nlwp
    rss=$(awk '/^VmRSS:/ {print $2}' "/proc/$pid/status" 2>/dev/null || echo 0)
    vsz=$(awk '/^VmSize:/ {print $2}' "/proc/$pid/status" 2>/dev/null || echo 0)
    nlwp=$(awk '/^Threads:/ {print $2}' "/proc/$pid/status" 2>/dev/null || echo 0)
    if [ "$have_pidstat" -eq 1 ]; then
      cpu=$(pidstat -p "$pid" 1 1 2>/dev/null | awk -v p="$pid" '$1 ~ /^[0-9]/ && $3==p {print $8; exit}')
      cpu=${cpu:-0.00}
    else
      cpu=0.00
    fi
    local stk=0 heap=0

      local sys_cpu_pct=0.00
      if [ -r /proc/stat ]; then
        read -r _ c_user c_nice c_system c_idle c_iowait c_irq c_softirq c_steal _ _ < /proc/stat
        local cpu_total=$((c_user + c_nice + c_system + c_idle + c_iowait + c_irq + c_softirq + c_steal))
        local cpu_idle=$((c_idle + c_iowait))
        local dt_total=$((cpu_total - prev_cpu_total))
        local dt_idle=$((cpu_idle - prev_cpu_idle))
        sys_cpu_pct=$(awk -v t="$dt_total" -v i="$dt_idle" 'BEGIN{ if (t<=0) printf "0.00"; else printf "%.2f", (100.0*(t-i))/t }')
        prev_cpu_total=$cpu_total
        prev_cpu_idle=$cpu_idle
      fi

      printf "%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s\n" \
        "$(date -Iseconds)" "$elapsed_s" "$load1" "$load5" "$load15" \
        "$mem_total" "$mem_avail" "$mem_used" "$proc_count" "$thread_total" \
        "$pid" "$cpu" "$sys_cpu_pct" "${rss:-0}" "${vsz:-0}" "${stk:-0}" "${heap:-0}" "${nlwp:-0}" \
        >> "$csv"
    fi

    if [ "$proc_step" -gt 0 ]; then
      sleep 0.1
    else
      sleep 1
    fi
  done
}

run_with_sampling() {
  local name="$1"
  local cmd="$2"
  local csv="$3"
  local max_seconds="$4"
  local proc_step="${5:-0}"

  log "Starting $name: $cmd"

  echo "timestamp,elapsed_s,load1,load5,load15,mem_total_kb,mem_available_kb,mem_used_kb,proc_count,thread_count_total,pid,cpu_pct,sys_cpu_pct,rss_kb,vsz_kb,stack_kb,heap_kb,proc_threads" > "$csv"

  local start_uptime pid
  start_uptime=$(awk '{print $1}' /proc/uptime)

  if command -v setsid >/dev/null 2>&1; then
    setsid bash -c "$cmd" &
    pid=$!
  else
    bash -c "$cmd" &
    pid=$!
  fi

  sample_process_csv "$pid" "$csv" "$start_uptime" "$max_seconds" "$proc_step"

  set +e
  wait "$pid"
  local exit_code=$?
  set -e

  log "Finished $name (exit_code=$exit_code)"
}

compile_creation_time() {
  log "Compiling creation_time"
  gcc -O2 -pthread -o "$ROOT_DIR/creation_time_experiment/creation_time" \
    "$ROOT_DIR/creation_time_experiment/creation_time.c"
}

compile_fork_bomb() {
  log "Compiling fork_bomb"
  gcc -O2 -o "$ROOT_DIR/fork_bomb/fork_bomb" \
    "$ROOT_DIR/fork_bomb/fork_bomb.c"
}

compile_thread_stress() {
  log "Compiling thread_stress"
  gcc -O2 -pthread -o "$ROOT_DIR/thread_stress_test/thread_stress" \
    "$ROOT_DIR/thread_stress_test/thread_stress.c"
}

compile_creation_time
compile_fork_bomb
compile_thread_stress

# Defaults (override by env vars)
ALLOW_DANGEROUS=${ALLOW_DANGEROUS:-0}
FORK_BOMB_SECONDS=${FORK_BOMB_SECONDS:-5}
PROC_SAMPLE_STEP=${PROC_SAMPLE_STEP:-100}
CREATION_TIME_ITERS=${CREATION_TIME_ITERS:-"100 1000 10000"}
CREATION_TIME_SECONDS=${CREATION_TIME_SECONDS:-60}

CREATION_TIME_STDOUT="$OUT_DIR/creation_time_stdout.log"
CREATION_TIME_STDERR="$OUT_DIR/creation_time_stderr.log"

run_with_sampling \
  "creation_time" \
  "cd '$ROOT_DIR/creation_time_experiment' && { for n in $CREATION_TIME_ITERS; do echo \"=== iterations=\$n ===\"; ./creation_time -n \"\$n\"; done; } >'$CREATION_TIME_STDOUT' 2>'$CREATION_TIME_STDERR'" \
  "$OUT_DIR/creation_time.csv" \
  "$CREATION_TIME_SECONDS" \
  0

if [ "$ALLOW_DANGEROUS" -eq 1 ]; then
  run_with_sampling \
    "fork_bomb" \
    "cd '$ROOT_DIR/fork_bomb' && ./fork_bomb" \
    "$OUT_DIR/fork_bomb.csv" \
    "$FORK_BOMB_SECONDS" \
    "$PROC_SAMPLE_STEP"
else
  log "Skipping fork_bomb (set ALLOW_DANGEROUS=1 to run)"
  echo "timestamp,elapsed_s,load1,load5,load15,mem_total_kb,mem_available_kb,mem_used_kb,proc_count,thread_count_total,pid,cpu_pct,sys_cpu_pct,rss_kb,vsz_kb,stack_kb,heap_kb,proc_threads" > "$OUT_DIR/fork_bomb.csv"
fi

run_with_sampling \
  "thread_recursive_scaling" \
  "cd '$ROOT_DIR/thread_recursive_scaling' && ./run_thread_recursive.sh 6 6 6" \
  "$OUT_DIR/thread_recursive_scaling.csv" \
  0 \
  0

run_with_sampling \
  "thread_stress" \
  "cd '$ROOT_DIR/thread_stress_test' && ./thread_stress" \
  "$OUT_DIR/thread_stress.csv" \
  0 \
  0

log "All experiments finished"
log "Run directory: $OUT_DIR"