#!/usr/bin/env bash
set -uo pipefail

# benchmark.sh — Performance benchmarking for tunectl
# Uses sysbench and fio to measure CPU, memory, and disk I/O performance.
#
# Modes:
#   (default)   Run all benchmarks and display current results
#   --baseline  Run benchmarks and save results for later comparison
#   --compare   Run benchmarks and compare against saved baseline
#
# Environment:
#   BENCHMARK_DIR  Override results directory (default: /var/lib/tunectl/benchmarks)
#
# Exit codes: 0=success, 1=operational failure, 2=usage error

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Results directory — can be overridden via environment for testing
BENCHMARK_DIR="${BENCHMARK_DIR:-/var/lib/tunectl/benchmarks}"
BASELINE_FILE="$BENCHMARK_DIR/baseline.json"

# Bounded resource parameters (safe for small systems, VAL-BENCH-017)
CPU_TEST_TIME=10          # seconds per CPU test (≤30s)
MEMORY_TEST_TIME=10       # seconds per memory test (≤30s)
FIO_RUNTIME=10            # seconds for fio test (≤30s)
FIO_SIZE="64M"            # fio file size (≤256MB)
FIO_IODEPTH=16            # fio queue depth (≤64)

# Temp files to track for cleanup
CLEANUP_FILES=()

# -------------------------------------------------------
# Cleanup handler — remove all temp files (VAL-BENCH-016)
# -------------------------------------------------------
cleanup() {
  for f in "${CLEANUP_FILES[@]}"; do
    rm -f "$f" 2>/dev/null
  done
}
trap cleanup EXIT

# -------------------------------------------------------
# Usage
# -------------------------------------------------------
usage() {
  cat <<EOF
Usage: benchmark.sh [OPTIONS]

Performance benchmarking using sysbench and fio.

Options:
  --baseline    Run benchmarks and save results for later comparison
  --compare     Run benchmarks and compare against saved baseline
  --help        Show this help message

Modes:
  (default)     Run benchmarks and display current results
  --baseline    Save results to $BENCHMARK_DIR/
  --compare     Load saved baseline and show comparison table

Tests:
  CPU single-thread     sysbench cpu --threads=1
  CPU multi-thread      sysbench cpu --threads=N (auto-detected)
  Memory read           sysbench memory --memory-oper=read
  Memory write          sysbench memory --memory-oper=write
  Disk random read      fio --rw=randread --bs=4k on disk-backed storage
EOF
  exit 0
}

# -------------------------------------------------------
# Parse arguments
# -------------------------------------------------------
MODE="default"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --baseline)
      MODE="baseline"
      shift
      ;;
    --compare)
      MODE="compare"
      shift
      ;;
    --help)
      usage
      ;;
    *)
      echo "Error: Unknown argument '$1'" >&2
      echo "Usage: benchmark.sh [--baseline|--compare|--help]" >&2
      exit 2
      ;;
  esac
done

# -------------------------------------------------------
# Check dependencies — graceful degradation
# Sets HAS_SYSBENCH and HAS_FIO globals (0=available, 1=missing)
# Exits 1 only if BOTH tools are missing
# -------------------------------------------------------
HAS_SYSBENCH=1
HAS_FIO=1

check_dependencies() {
  local missing=()

  if command -v sysbench &>/dev/null; then
    HAS_SYSBENCH=0
  else
    missing+=("sysbench")
    echo "▲  sysbench not found. Install: sudo apt install sysbench" >&2
  fi

  if command -v fio &>/dev/null; then
    HAS_FIO=0
  else
    missing+=("fio")
    echo "▲  fio not found. Install: sudo apt install fio" >&2
  fi

  if [[ ${#missing[@]} -eq 2 ]]; then
    echo "" >&2
    echo "Error: No benchmark tools available. Install at least one:" >&2
    echo "  sudo apt install sysbench fio" >&2
    exit 1
  fi
}

# -------------------------------------------------------
# Find a disk-backed path for fio tests (NOT tmpfs)
# VAL-BENCH-006: Must use actual block storage
# -------------------------------------------------------
find_disk_backed_path() {
  # Check common paths, prefer /var/tmp (usually disk-backed) or /home
  local candidates=("/var/tmp" "/home" "/var/lib" "/opt")

  for candidate in "${candidates[@]}"; do
    if [[ -d "$candidate" && -w "$candidate" ]]; then
      local fstype
      fstype=$(stat -f -c '%T' "$candidate" 2>/dev/null || findmnt -n -o FSTYPE --target "$candidate" 2>/dev/null || echo "unknown")
      # tmpfs has a specific magic number 0x01021994 which stat -f shows as "tmpfs"
      if [[ "$fstype" != "tmpfs" ]]; then
        echo "$candidate"
        return 0
      fi
    fi
  done

  # Fallback: use /var/tmp even if we can't detect type
  if [[ -d "/var/tmp" && -w "/var/tmp" ]]; then
    echo "/var/tmp"
    return 0
  fi

  echo "Error: Cannot find a disk-backed writable directory for fio tests" >&2
  return 1
}

# -------------------------------------------------------
# CPU benchmark — single thread
# -------------------------------------------------------
run_cpu_single() {
  local raw cmd_exit=0
  raw=$(sysbench cpu --threads=1 --time="$CPU_TEST_TIME" run 2>/dev/null) || cmd_exit=$?

  if [[ $cmd_exit -ne 0 || -z "$raw" ]]; then
    echo "Error: CPU single-thread benchmark failed (sysbench exit code $cmd_exit)" >&2
    echo "ERROR"
    return 0
  fi

  # Parse "events per second: NNN.NN"
  local eps
  eps=$(echo "$raw" | grep -oP 'events per second:\s*\K[0-9]+(\.[0-9]+)?' || true)

  if [[ -z "$eps" ]]; then
    echo "Error: CPU single-thread benchmark produced no parseable output" >&2
    echo "ERROR"
    return 0
  fi

  echo "$eps"
}

# -------------------------------------------------------
# CPU benchmark — multi thread
# -------------------------------------------------------
run_cpu_multi() {
  local threads
  threads=$(nproc 2>/dev/null || echo 1)
  local raw cmd_exit=0
  raw=$(sysbench cpu --threads="$threads" --time="$CPU_TEST_TIME" run 2>/dev/null) || cmd_exit=$?

  if [[ $cmd_exit -ne 0 || -z "$raw" ]]; then
    echo "Error: CPU multi-thread benchmark failed (sysbench exit code $cmd_exit)" >&2
    echo "ERROR"
    return 0
  fi

  local eps
  eps=$(echo "$raw" | grep -oP 'events per second:\s*\K[0-9]+(\.[0-9]+)?' || true)

  if [[ -z "$eps" ]]; then
    echo "Error: CPU multi-thread benchmark produced no parseable output" >&2
    echo "ERROR"
    return 0
  fi

  echo "$eps"
}

# -------------------------------------------------------
# Memory benchmark — read
# -------------------------------------------------------
run_mem_read() {
  local raw cmd_exit=0
  raw=$(sysbench memory --memory-oper=read --threads=1 --time="$MEMORY_TEST_TIME" run 2>/dev/null) || cmd_exit=$?

  if [[ $cmd_exit -ne 0 || -z "$raw" ]]; then
    echo "Error: Memory read benchmark failed (sysbench exit code $cmd_exit)" >&2
    echo "ERROR"
    return 0
  fi

  # Parse "NNNN.NN MiB transferred (NNNN.NN MiB/sec)"
  local mib_sec
  mib_sec=$(echo "$raw" | grep -oP '\(([0-9]+(\.[0-9]+)?)\s*MiB/sec\)' | grep -oP '[0-9]+(\.[0-9]+)?' | head -1 || true)

  if [[ -z "$mib_sec" ]]; then
    echo "Error: Memory read benchmark produced no parseable output" >&2
    echo "ERROR"
    return 0
  fi

  echo "$mib_sec"
}

# -------------------------------------------------------
# Memory benchmark — write
# -------------------------------------------------------
run_mem_write() {
  local raw cmd_exit=0
  raw=$(sysbench memory --memory-oper=write --threads=1 --time="$MEMORY_TEST_TIME" run 2>/dev/null) || cmd_exit=$?

  if [[ $cmd_exit -ne 0 || -z "$raw" ]]; then
    echo "Error: Memory write benchmark failed (sysbench exit code $cmd_exit)" >&2
    echo "ERROR"
    return 0
  fi

  local mib_sec
  mib_sec=$(echo "$raw" | grep -oP '\(([0-9]+(\.[0-9]+)?)\s*MiB/sec\)' | grep -oP '[0-9]+(\.[0-9]+)?' | head -1 || true)

  if [[ -z "$mib_sec" ]]; then
    echo "Error: Memory write benchmark produced no parseable output" >&2
    echo "ERROR"
    return 0
  fi

  echo "$mib_sec"
}

# -------------------------------------------------------
# Disk benchmark — random read IOPS
# -------------------------------------------------------
run_disk_rand_read() {
  local disk_path
  disk_path=$(find_disk_backed_path) || {
    echo "Error: Disk random read benchmark failed (no disk-backed path found)" >&2
    echo "ERROR"
    return 0
  }

  local fio_file="${disk_path}/benchmark_fio_testfile"
  CLEANUP_FILES+=("$fio_file")

  local raw cmd_exit=0
  raw=$(fio --name=benchmark_fio_test \
    --filename="$fio_file" \
    --rw=randread \
    --bs=4k \
    --direct=1 \
    --ioengine=libaio \
    --iodepth="$FIO_IODEPTH" \
    --size="$FIO_SIZE" \
    --runtime="$FIO_RUNTIME" \
    --time_based \
    --output-format=json 2>/dev/null) || cmd_exit=$?

  # Cleanup the fio test file immediately
  rm -f "$fio_file" 2>/dev/null

  if [[ $cmd_exit -ne 0 || -z "$raw" ]]; then
    echo "Error: Disk random read benchmark failed (fio exit code $cmd_exit)" >&2
    echo "ERROR"
    return 0
  fi

  # Parse IOPS from JSON output
  local iops
  iops=$(echo "$raw" | python3 -c "
import json, sys
try:
    d = json.load(sys.stdin)
    print(f\"{d['jobs'][0]['read']['iops']:.2f}\")
except Exception:
    print('')
" 2>/dev/null)

  if [[ -z "$iops" ]]; then
    echo "Error: Disk random read benchmark produced no parseable IOPS output" >&2
    echo "ERROR"
    return 0
  fi

  echo "$iops"
}

# -------------------------------------------------------
# Run all benchmarks and collect results
# Returns results as key=value pairs
# Skips benchmarks for missing tools (uses HAS_SYSBENCH/HAS_FIO)
# Exits 1 if any executed benchmark command failed (ERROR sentinel)
# -------------------------------------------------------
run_all_benchmarks() {
  local failed_benchmarks=()
  local step=0 total=0

  # Count total benchmarks to run
  [[ $HAS_SYSBENCH -eq 0 ]] && total=$((total + 4))
  [[ $HAS_FIO -eq 0 ]] && total=$((total + 1))

  echo "Running benchmarks..." >&2

  local cpu_st="SKIPPED" cpu_mt="SKIPPED" mem_read="SKIPPED" mem_write="SKIPPED" disk_rr="SKIPPED"

  if [[ $HAS_SYSBENCH -eq 0 ]]; then
    step=$((step + 1))
    echo "  [$step/$total] CPU single-thread..." >&2
    cpu_st=$(run_cpu_single)
    [[ "$cpu_st" == "ERROR" ]] && failed_benchmarks+=("CPU single-thread")

    step=$((step + 1))
    echo "  [$step/$total] CPU multi-thread..." >&2
    cpu_mt=$(run_cpu_multi)
    [[ "$cpu_mt" == "ERROR" ]] && failed_benchmarks+=("CPU multi-thread")

    step=$((step + 1))
    echo "  [$step/$total] Memory read..." >&2
    mem_read=$(run_mem_read)
    [[ "$mem_read" == "ERROR" ]] && failed_benchmarks+=("Memory read")

    step=$((step + 1))
    echo "  [$step/$total] Memory write..." >&2
    mem_write=$(run_mem_write)
    [[ "$mem_write" == "ERROR" ]] && failed_benchmarks+=("Memory write")
  else
    echo "  Skipping sysbench tests (sysbench not installed)" >&2
  fi

  if [[ $HAS_FIO -eq 0 ]]; then
    step=$((step + 1))
    echo "  [$step/$total] Disk random read..." >&2
    disk_rr=$(run_disk_rand_read)
    [[ "$disk_rr" == "ERROR" ]] && failed_benchmarks+=("Disk random read")
  else
    echo "  Skipping fio tests (fio not installed)" >&2
  fi

  echo "Done." >&2

  # Output results as structured key-value lines
  cat <<EOF
cpu_single_thread $cpu_st events/s
cpu_multi_thread $cpu_mt events/s
mem_read $mem_read MiB/s
mem_write $mem_write MiB/s
disk_rand_read_iops $disk_rr IOPS
EOF

  # Check if any benchmarks failed
  if [[ ${#failed_benchmarks[@]} -gt 0 ]]; then
    echo "" >&2
    echo "Error: The following benchmark(s) failed:" >&2
    for b in "${failed_benchmarks[@]}"; do
      echo "  - $b" >&2
    done
    return 1
  fi

  return 0
}

# -------------------------------------------------------
# Display results in structured format (VAL-BENCH-011)
# -------------------------------------------------------
display_results() {
  local results="$1"

  echo ""
  echo "===== tunectl Benchmark Results ====="
  echo ""
  printf "  %-30s %15s %s\n" "Metric" "Value" "Unit"
  printf "  %-30s %15s %s\n" "------------------------------" "---------------" "----"

  while IFS=' ' read -r metric value unit; do
    [[ -z "$metric" ]] && continue
    local label
    case "$metric" in
      cpu_single_thread)     label="CPU Single-Thread" ;;
      cpu_multi_thread)      label="CPU Multi-Thread" ;;
      mem_read)              label="Memory Read" ;;
      mem_write)             label="Memory Write" ;;
      disk_rand_read_iops)   label="Disk Random Read" ;;
      *)                     label="$metric" ;;
    esac
    # Show N/A for ERROR values, show skipped for SKIPPED
    if [[ "$value" == "SKIPPED" ]]; then
      printf "  %-30s %15s %s\n" "$label" "N/A" "(skipped)"
    elif [[ "$value" == "ERROR" ]]; then
      printf "  %-30s %15s %s\n" "$label" "N/A" "(failed)"
    else
      printf "  %-30s %15s %s\n" "$label" "$value" "$unit"
    fi
  done <<< "$results"

  echo ""
  echo "====================================="
}

# -------------------------------------------------------
# Save results as JSON (for --baseline mode)
# -------------------------------------------------------
save_results_json() {
  local results="$1"
  local outfile="$2"

  # Ensure output directory exists
  local outdir
  outdir=$(dirname "$outfile")
  mkdir -p "$outdir" 2>/dev/null || {
    echo "Error: Cannot create directory $outdir" >&2
    return 1
  }

  # Build JSON — skip ERROR metrics (don't save bogus values)
  python3 -c "
import json, sys, datetime

results = {}
for line in sys.stdin:
    parts = line.strip().split(None, 2)
    if len(parts) >= 2:
        metric = parts[0]
        if parts[1] in ('ERROR', 'SKIPPED'):
            results[metric] = {'value': None, 'unit': parts[1]}
            continue
        value = float(parts[1])
        unit = parts[2] if len(parts) > 2 else ''
        results[metric] = {'value': value, 'unit': unit}

output = {
    'timestamp': datetime.datetime.now().isoformat(),
    'metrics': results
}

with open('$outfile', 'w') as f:
    json.dump(output, f, indent=2)
" <<< "$results"
}

# -------------------------------------------------------
# Load baseline JSON
# -------------------------------------------------------
load_baseline() {
  local file="$1"
  if [[ ! -f "$file" ]]; then
    return 1
  fi
  python3 -c "
import json, sys
try:
    d = json.load(open('$file'))
    for metric, data in d.get('metrics', d).items():
        if isinstance(data, dict):
            val = data['value']
            unit = data.get('unit', '')
            # Handle ERROR metrics saved with None value
            if val is None or unit in ('ERROR', 'SKIPPED'):
                print(f'{metric} {unit} {unit}')
            else:
                print(f'{metric} {val} {unit}')
        else:
            print(f'{metric} {data}')
except Exception as e:
    print(f'Error loading baseline: {e}', file=sys.stderr)
    sys.exit(1)
" 2>/dev/null
}

# -------------------------------------------------------
# Display comparison table (VAL-BENCH-003)
# -------------------------------------------------------
display_comparison() {
  local baseline_results="$1"
  local current_results="$2"

  echo ""
  echo "===== tunectl Benchmark Comparison ====="
  echo ""
  printf "  %-25s %12s %12s %12s %8s\n" "Metric" "Baseline" "Current" "Change" "%Change"
  printf "  %-25s %12s %12s %12s %8s\n" "-------------------------" "------------" "------------" "------------" "--------"

  # Build associative arrays
  declare -A baseline_vals
  declare -A baseline_units
  while IFS=' ' read -r metric value unit; do
    [[ -z "$metric" ]] && continue
    baseline_vals["$metric"]="$value"
    baseline_units["$metric"]="$unit"
  done <<< "$baseline_results"

  while IFS=' ' read -r metric value unit; do
    [[ -z "$metric" ]] && continue

    local label
    case "$metric" in
      cpu_single_thread)     label="CPU Single-Thread" ;;
      cpu_multi_thread)      label="CPU Multi-Thread" ;;
      mem_read)              label="Memory Read" ;;
      mem_write)             label="Memory Write" ;;
      disk_rand_read_iops)   label="Disk Random Read" ;;
      *)                     label="$metric" ;;
    esac

    # Handle SKIPPED or ERROR in current value — show N/A instead of bogus computation
    if [[ "$value" == "SKIPPED" || "$value" == "ERROR" ]]; then
      local bval="${baseline_vals[$metric]:-N/A}"
      printf "  %-25s %12s %12s %12s %8s\n" "$label" "$bval" "N/A" "N/A" "N/A"
      continue
    fi

    local bval="${baseline_vals[$metric]:-}"
    # Handle missing, ERROR, or SKIPPED baseline — show N/A for comparison
    if [[ -z "$bval" || "$bval" == "None" || "$bval" == "ERROR" || "$bval" == "SKIPPED" ]]; then
      printf "  %-25s %12s %12s %12s %8s\n" "$label" "N/A" "$value" "N/A" "N/A"
      continue
    fi

    # Calculate change and percentage
    local change pct sign
    change=$(python3 -c "b=$bval; c=$value; print(f'{c-b:.2f}')" 2>/dev/null || echo "0")
    pct=$(python3 -c "
b=$bval; c=$value
if b == 0:
    print('N/A')
else:
    pct = ((c - b) / b) * 100
    print(f'{pct:+.1f}%')
" 2>/dev/null || echo "N/A")

    # Add sign to change for display
    sign=$(python3 -c "c=$change; print('+' if c > 0 else '')" 2>/dev/null || echo "")
    printf "  %-25s %12s %12s %12s %8s\n" "$label" "$bval" "$value" "${sign}${change}" "$pct"

  done <<< "$current_results"

  echo ""
  echo "========================================="
}

# -------------------------------------------------------
# Main
# -------------------------------------------------------
main() {
  check_dependencies

  case "$MODE" in
    default)
      local results benchmark_rc=0
      results=$(run_all_benchmarks) || benchmark_rc=$?
      display_results "$results"
      if [[ $benchmark_rc -ne 0 ]]; then
        exit 1
      fi
      ;;

    baseline)
      local results benchmark_rc=0
      results=$(run_all_benchmarks) || benchmark_rc=$?
      display_results "$results"

      if [[ $benchmark_rc -ne 0 ]]; then
        echo "Warning: Some benchmarks failed. Saving partial results." >&2
      fi

      save_results_json "$results" "$BASELINE_FILE" || {
        echo "Error: Failed to save baseline results" >&2
        exit 1
      }
      echo "Baseline saved to: $BASELINE_FILE"

      if [[ $benchmark_rc -ne 0 ]]; then
        exit 1
      fi
      ;;

    compare)
      if [[ ! -f "$BASELINE_FILE" ]]; then
        echo "Error: No baseline found at $BASELINE_FILE" >&2
        echo "Run 'benchmark.sh --baseline' first to capture a baseline." >&2
        exit 1
      fi

      local baseline_results
      baseline_results=$(load_baseline "$BASELINE_FILE") || {
        echo "Error: Failed to load baseline from $BASELINE_FILE" >&2
        exit 1
      }

      local current_results benchmark_rc=0
      current_results=$(run_all_benchmarks) || benchmark_rc=$?

      display_comparison "$baseline_results" "$current_results"

      if [[ $benchmark_rc -ne 0 ]]; then
        exit 1
      fi
      ;;
  esac
}

main