#!/usr/bin/env bash
#
# xz-benchmark.sh — Run the xz single-thread and multi-thread benchmarks
# on a ~1GB text file, collect CPU info, and output a pasteable block.
#
# Usage:
#   ./xz-benchmark.sh [enwik9] [submitter]
#
# If enwik9 is not present, it will be downloaded automatically.
# The submitter name defaults to $USER.
#
set -euo pipefail

# ── Configuration ─────────────────────────────────────────────────────────

INPUT_FILE="${1:-enwik9}"
SUBMITTER="${2:-$USER}"
BENCH_DIR="$(cd "$(dirname "$0")" && pwd)"
OUTPUT_FILE="${INPUT_FILE}.xz"

# ── Helpers ───────────────────────────────────────────────────────────────

err()  { echo "ERROR: $*" >&2; exit 1; }
info() { echo ">>> $*" >&2; }

# Parse "real Xm Y.YYYs" from bash time output (TIMEFORMAT)
parse_real() {
    local line="$1"
    # Accept formats like "1m50.963s" or "110.963s"
    if [[ "$line" =~ ([0-9]+)m([0-9.]+)s ]]; then
        local min="${BASH_REMATCH[1]}"
        local sec="${BASH_REMATCH[2]}"
        printf "%dm%ss\n" "$min" "$sec"
    elif [[ "$line" =~ ([0-9.]+)s ]]; then
        printf "0m%ss\n" "${BASH_REMATCH[1]}"
    else
        printf "%s\n" "$line"
    fi
}

# ── Prerequisites ─────────────────────────────────────────────────────────

command -v xz >/dev/null 2>&1 || err "xz not found in PATH"
command -v lscpu >/dev/null 2>&1 || err "lscpu not found (install util-linux)"

# ── Download input file if needed ─────────────────────────────────────────

if [[ ! -f "$INPUT_FILE" ]]; then
    info "Downloading enwik9 (~1GB) from https://cs.fit.edu/~mmahoney/compression/"
    if command -v wget >/dev/null 2>&1; then
        wget -q -O "$INPUT_FILE" "https://cs.fit.edu/~mmahoney/compression/enwik9.zip"
    elif command -v curl >/dev/null 2>&1; then
        curl -sL -o "$INPUT_FILE" "https://cs.fit.edu/~mmahoney/compression/enwik9.zip"
    else
        err "Need wget or curl to download enwik9"
    fi

    # enwik9.zip contains enwik9
    if file "$INPUT_FILE" | grep -qi zip; then
        info "Unzipping..."
        unzip -o "$INPUT_FILE" && rm -f "$INPUT_FILE" && mv enwik9 "$INPUT_FILE" 2>/dev/null || true
    fi

    [[ -f "$INPUT_FILE" ]] || err "Failed to obtain $INPUT_FILE"
fi

info "Input file: $INPUT_FILE ($(du -h "$INPUT_FILE" | cut -f1))"

# ── Collect system info ───────────────────────────────────────────────────

info "Collecting CPU information..."

# CPU model name
CPU_MODEL="$(lscpu | grep -m1 'Model name:' | sed 's/Model name:\s*//; s/  */ /g' | xargs)"

# CPU count info
CPU_SOCKETS="$(lscpu | grep -m1 'Socket(s):' | awk '{print $2}')"
CORES_PER_SOCKET="$(lscpu | grep -m1 'Core(s) per socket:' | awk '{print $4}')"
THREADS_PER_CORE="$(lscpu | grep -m1 'Thread(s) per core:' | awk '{print $4}')"

TOTAL_CORES=$(( CPU_SOCKETS * CORES_PER_SOCKET ))
TOTAL_THREADS=$(( TOTAL_CORES * THREADS_PER_CORE ))

# Clock speeds
CPU_MHZ="$(lscpu | grep -m1 'CPU max MHz:' | awk -F: '{gsub(/^ +| +$/,"",$2); print $2}')"
if [[ -z "$CPU_MHZ" || "$CPU_MHZ" == *"-"* ]]; then
    CPU_MHZ="$(lscpu | grep -m1 'CPU MHz:' | awk -F: '{gsub(/^ +| +$/,"",$2); print $2}')"
fi

# Base and boost clocks (best effort)
BASE_CLOCK="$(lscpu | grep -m1 'CPU min MHz:' | awk -F: '{gsub(/^ +| +$/,"",$2); print $2}')"
BOOST_CLOCK="$CPU_MHZ"

if [[ -n "$BASE_CLOCK" && "$BASE_CLOCK" != *"-"* ]]; then
    BASE_CLOCK="$(awk "BEGIN{printf \"%.1f\", $BASE_CLOCK / 1000}") GHz"
else
    BASE_CLOCK="-"
fi

if [[ -n "$BOOST_CLOCK" && "$BOOST_CLOCK" != *"-"* ]]; then
    BOOST_CLOCK="$(awk "BEGIN{printf \"%.1f\", $BOOST_CLOCK / 1000}") GHz"
else
    BOOST_CLOCK="-"
fi

# Architecture
ARCH="$(uname -m)"

# RAM
RAM_KB="$(grep MemTotal /proc/meminfo | awk '{print $2}')"
RAM_GB="$(awk "BEGIN{printf \"%.0f\", $RAM_KB / 1048576}")"
RAM="${RAM_GB} GB"

# OS
if [[ -f /etc/os-release ]]; then
    . /etc/os-release
    OS="${PRETTY_NAME:-$(uname -sr)}"
else
    OS="$(uname -sr)"
fi

# Hostname
HOST="$(hostname -s 2>/dev/null || hostname)"

# xz version
XZ_VERSION="$(xz --version | head -1)"

info "Host: $HOST | CPU: $CPU_MODEL | Cores: $TOTAL_CORES | Arch: $ARCH"

# ── Run benchmarks ────────────────────────────────────────────────────────

# We use bash's TIMEFORMAT to capture times in a parseable format
# TIMEFORMAT='%PR %PU %PS' gives real, user, sys in seconds
# But we want the "Xm Y.YYYs" format to match existing data.
# We'll use the `time` builtin with a custom format, then convert.

run_bench() {
    local threads="$1"
    local label="$2"
    local time_real time_user time_sys
    local output

    info "Running xz -T$threads ($label)..."

    # Use TIMEFORMAT to get seconds, then convert to m/s format
    # Capturing time output requires a subshell trick
    {
        TIMEFORMAT=$'%R %U %S'
        time xz -k -T"$threads" "$INPUT_FILE"
    } 2> "$BENCH_DIR/.bench_time_tmp"

    read -r time_real time_user time_sys < "$BENCH_DIR/.bench_time_tmp"
    rm -f "$BENCH_DIR/.bench_time_tmp" "$OUTPUT_FILE"

    # Convert seconds → "Xm Y.YYYs" format
    local real_fmt user_fmt sys_fmt
    real_fmt=$(sec_to_fmt "$time_real")
    user_fmt=$(sec_to_fmt "$time_user")
    sys_fmt=$(sec_to_fmt "$time_sys")

    echo "$real_fmt|$user_fmt|$sys_fmt"
}

sec_to_fmt() {
    local sec="$1"
    local int_part="${sec%.*}"
    [[ "$int_part" == "$sec" ]] && int_part="$sec"  # no decimal point
    local min=$(( int_part / 60 ))
    local rem
    rem=$(awk "BEGIN{printf \"%.3f\", $sec - ($min * 60)}")
    printf "%dm%ss\n" "$min" "$rem"
}

# Convert "Xm Y.YYYs" or seconds to a numeric value
sec_to_num() {
    local val="$1"
    if [[ "$val" =~ ([0-9]+)m([0-9.]+)s ]]; then
        awk "BEGIN{printf \"%.3f\", ${BASH_REMATCH[1]} * 60 + ${BASH_REMATCH[2]}}"
    elif [[ "$val" =~ ([0-9.]+)s ]]; then
        awk "BEGIN{printf \"%.3f\", ${BASH_REMATCH[1]}}"
    else
        printf "NULL"
    fi
}

# Run single-threaded benchmark
SYNC_RESULT=$(run_bench 1 "single-thread")
sleep 2
sync

ST_REAL="$(echo "$SYNC_RESULT" | cut -d'|' -f1)"
ST_USER="$(echo "$SYNC_RESULT" | cut -d'|' -f2)"
ST_SYS="$(echo "$SYNC_RESULT" | cut -d'|' -f3)"

# Run multi-threaded benchmark
MT_RESULT=$(run_bench 0 "multi-thread")
sync

MT_REAL="$(echo "$MT_RESULT" | cut -d'|' -f1)"
MT_USER="$(echo "$MT_RESULT" | cut -d'|' -f2)"
MT_SYS="$(echo "$MT_RESULT" | cut -d'|' -f3)"

# ── Output pasteable block ────────────────────────────────────────────────

info "Benchmark complete!"

# Cores string for multithread
if [[ "$CPU_SOCKETS" -gt 1 ]]; then
    MT_CORES="${CPU_SOCKETS}x${CORES_PER_SOCKET}"
    SMT_CORES="$TOTAL_THREADS"
else
    MT_CORES="$TOTAL_CORES"
    SMT_CORES="$TOTAL_THREADS"
fi

# Pre-compute numeric seconds for SQL
ST_REAL_NUM=$(sec_to_num "$ST_REAL")
ST_USER_NUM=$(sec_to_num "$ST_USER")
ST_SYS_NUM=$(sec_to_num "$ST_SYS")
MT_REAL_NUM=$(sec_to_num "$MT_REAL")
MT_USER_NUM=$(sec_to_num "$MT_USER")
MT_SYS_NUM=$(sec_to_num "$MT_SYS")

cat <<EOF

╔══════════════════════════════════════════════════════════════════════════╗
║                    XZ BENCHMARK RESULTS                                 ║
╠══════════════════════════════════════════════════════════════════════════╣

── System Information ──
  Host:         $HOST
  Model:        [enter model]
  Submitter:    $SUBMITTER
  CPU:          $CPU_MODEL
  Base Clock:   $BASE_CLOCK
  Boost Clock:  $BOOST_CLOCK
  Sockets:      $CPU_SOCKETS
  Cores/Socket: $CORES_PER_SOCKET
  Threads/Core: $THREADS_PER_CORE
  Total Cores:  $TOTAL_CORES
  Total Threads:$TOTAL_THREADS
  Arch:         $ARCH
  RAM:          $RAM
  OS:           $OS
  xz Version:   $XZ_VERSION
  Input File:   $INPUT_FILE ($(du -h "$INPUT_FILE" | cut -f1))

── Results ──

  Single Thread (-T1):
    real:  $ST_REAL
    user:  $ST_USER
    sys:   $ST_SYS

  Multi Thread (-T0):
    real:  $MT_REAL
    user:  $MT_USER
    sys:   $MT_SYS

── Pasteable Wiki Table Row (MediaWiki) ──

=== Single Thread ===
|-
! $HOST
| [enter model]
| $SUBMITTER
| $CPU_MODEL
| $BASE_CLOCK
| $BOOST_CLOCK
| $ARCH
| $RAM
| $OS
| $ST_REAL
| $ST_USER
| $ST_SYS

=== Multithread (Real Cores) ===
|-
! $HOST
| [enter model]
| $SUBMITTER
| $CPU_MODEL
| $BASE_CLOCK
| $BOOST_CLOCK
| $MT_CORES
| $ARCH
| $RAM
| $OS
| $MT_REAL
| $MT_USER
| $MT_SYS

── Pasteable Wiki Table Row (DokuWiki) ──

^ Host ^ Model ^ Submitter ^ CPU ^ Base Clock ^ Boost Clock ^ Cores ^ Arch ^ RAM ^ OS ^ real ^ user ^ sys ^
| $HOST | [enter model] | $SUBMITTER | $CPU_MODEL | $BASE_CLOCK | $BOOST_CLOCK | $MT_CORES | $ARCH | $RAM | $OS | $ST_REAL | $ST_USER | $ST_SYS |

── SQL INSERT (for direct database submission) ──

INSERT INTO results (test_type, host, model, submitter, cpu, base_clock, boost_clock, cores, arch, ram, os, real_time, real_time_sec, user_time, user_time_sec, sys_time, sys_time_sec)
VALUES
  ('single',      '$HOST', '[enter model]', '$SUBMITTER', '$CPU_MODEL', '$BASE_CLOCK', '$BOOST_CLOCK', NULL,        '$ARCH', '$RAM', '$OS', '$ST_REAL', $ST_REAL_NUM, '$ST_USER', $ST_USER_NUM, '$ST_SYS', $ST_SYS_NUM),
  ('multithread', '$HOST', '[enter model]', '$SUBMITTER', '$CPU_MODEL', '$BASE_CLOCK', '$BOOST_CLOCK', '$MT_CORES', '$ARCH', '$RAM', '$OS', '$MT_REAL', $MT_REAL_NUM, '$MT_USER', $MT_USER_NUM, '$MT_SYS', $MT_SYS_NUM);

── JSON (for DokuWiki plugin / API submission) ──

[
  {
    "test_type": "single",
    "host": "$HOST",
    "model": "[enter model]",
    "submitter": "$SUBMITTER",
    "cpu": "$CPU_MODEL",
    "base_clock": "$BASE_CLOCK",
    "boost_clock": "$BOOST_CLOCK",
    "cores": null,
    "arch": "$ARCH",
    "ram": "$RAM",
    "os": "$OS",
    "real_time": "$ST_REAL",
    "user_time": "$ST_USER",
    "sys_time": "$ST_SYS"
  },
  {
    "test_type": "multithread",
    "host": "$HOST",
    "model": "[enter model]",
    "submitter": "$SUBMITTER",
    "cpu": "$CPU_MODEL",
    "base_clock": "$BASE_CLOCK",
    "boost_clock": "$BOOST_CLOCK",
    "cores": "$MT_CORES",
    "arch": "$ARCH",
    "ram": "$RAM",
    "os": "$OS",
    "real_time": "$MT_REAL",
    "user_time": "$MT_USER",
    "sys_time": "$MT_SYS"
  }
]

╚══════════════════════════════════════════════════════════════════════════╝
EOF