#!/bin/bash
# bench-boot.sh — Measure KubeSolo OS boot performance in QEMU
#
# Measures:
#   - Time to first console output (kernel loaded)
#   - Time to init complete (all stages done)
#   - Time to K8s node Ready
#   - Time to first pod Running (nginx test)
#   - Peak memory usage
#   - Disk image/ISO size
#
# Usage:
#   test/benchmark/bench-boot.sh <iso-or-img> [--runs N]
#
# Output: JSON benchmark results to stdout, human-readable to stderr
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
IMAGE="${1:?Usage: bench-boot.sh <iso-or-img> [--runs N]}"
RUNS=3
SSH_PORT=2222
K8S_PORT=6443

. "$SCRIPT_DIR/../lib/qemu-helpers.sh"

shift || true
while [ $# -gt 0 ]; do
    case "$1" in
        --runs) RUNS="$2"; shift 2 ;;
        *) echo "Unknown option: $1" >&2; exit 1 ;;
    esac
done

if [ ! -f "$IMAGE" ]; then
    echo "ERROR: Image not found: $IMAGE" >&2
    exit 1
fi

# Determine image type
IMAGE_TYPE="iso"
if [[ "$IMAGE" == *.img ]]; then
    IMAGE_TYPE="disk"
fi

echo "=== KubeSolo OS Boot Benchmark ===" >&2
echo "Image:  $IMAGE ($(du -h "$IMAGE" | cut -f1))" >&2
echo "Type:   $IMAGE_TYPE" >&2
echo "Runs:   $RUNS" >&2
echo "" >&2

EXTRACT_DIR=""
TEMP_DISK=""

cleanup() {
    [ -n "$TEMP_DISK" ] && rm -f "$TEMP_DISK"
    [ -n "$EXTRACT_DIR" ] && rm -rf "$EXTRACT_DIR"
}
trap cleanup EXIT

# Build QEMU command
QEMU_CMD=(
    qemu-system-x86_64
    -m 1024
    -smp 2
    -nographic
    -no-reboot
    -serial mon:stdio
    -net "nic,model=virtio"
    -net "user,hostfwd=tcp::${SSH_PORT}-:22,hostfwd=tcp::${K8S_PORT}-:6443"
)

# Add KVM if available
KVM_FLAG=$(detect_kvm)
if [ -n "$KVM_FLAG" ]; then
    QEMU_CMD+=(-enable-kvm -cpu host)
    echo "KVM:    enabled" >&2
else
    QEMU_CMD+=(-cpu max)
    echo "KVM:    not available (TCG)" >&2
fi
echo "" >&2

if [ "$IMAGE_TYPE" = "iso" ]; then
    # Extract kernel/initramfs for direct boot (required for -append to work)
    EXTRACT_DIR="$(mktemp -d /tmp/kubesolo-bench-extract-XXXXXX)"
    extract_kernel_from_iso "$IMAGE" "$EXTRACT_DIR" >&2
    QEMU_CMD+=(-kernel "$VMLINUZ" -initrd "$INITRAMFS")
    QEMU_CMD+=(-append "console=ttyS0,115200n8 kubesolo.debug")
    # Add a temp disk for persistence
    TEMP_DISK=$(mktemp /tmp/kubesolo-bench-XXXXXX.img)
    qemu-img create -f qcow2 "$TEMP_DISK" 8G >/dev/null 2>&1
    QEMU_CMD+=(-drive "file=$TEMP_DISK,format=qcow2,if=virtio")
else
    QEMU_CMD+=(-drive "file=$IMAGE,format=raw,if=virtio")
fi

# Results arrays
declare -a BOOT_TIMES
declare -a INIT_TIMES
declare -a K8S_TIMES
declare -a MEMORY_USAGE

for run in $(seq 1 "$RUNS"); do
    echo "--- Run $run/$RUNS ---" >&2

    START_TIME=$(date +%s%N)
    BOOT_DONE=""
    INIT_DONE=""
    K8S_READY=""
    PEAK_MEM=""

    # Create a log file for this run
    LOG=$(mktemp /tmp/kubesolo-bench-log-XXXXXX)

    # Run QEMU with timeout, capturing output
    timeout 300 "${QEMU_CMD[@]}" 2>&1 | while IFS= read -r line; do
        NOW=$(date +%s%N)
        ELAPSED_MS=$(( (NOW - START_TIME) / 1000000 ))

        echo "$line" >> "$LOG"

        # Detect boot milestones from serial output
        case "$line" in
            *"Linux version"*)
                if [ -z "$BOOT_DONE" ]; then
                    BOOT_DONE="$ELAPSED_MS"
                    echo "  Kernel loaded: ${ELAPSED_MS}ms" >&2
                    echo "KERNEL_MS=$ELAPSED_MS" >> "$LOG.times"
                fi
                ;;
            *"KubeSolo is running"*|*"kubesolo-init"*"OK"*)
                if [ -z "$INIT_DONE" ]; then
                    INIT_DONE="$ELAPSED_MS"
                    echo "  Init complete: ${ELAPSED_MS}ms" >&2
                    echo "INIT_MS=$ELAPSED_MS" >> "$LOG.times"
                fi
                ;;
            *"node is Ready"*|*"NotReady"*"Ready"*|*"kubesolo"*"Ready"*)
                if [ -z "$K8S_READY" ]; then
                    K8S_READY="$ELAPSED_MS"
                    echo "  K8s Ready:     ${ELAPSED_MS}ms" >&2
                    echo "K8S_MS=$ELAPSED_MS" >> "$LOG.times"
                fi
                ;;
            *"MemTotal:"*|*"MemAvailable:"*)
                # Capture memory info if printed
                echo "MEM_LINE=$line" >> "$LOG.times"
                ;;
        esac

        # Stop after K8s is ready (or timeout)
        if [ -n "$K8S_READY" ]; then
            break
        fi
    done || true

    # Read results from log
    if [ -f "$LOG.times" ]; then
        KERNEL_MS=$(grep "KERNEL_MS=" "$LOG.times" 2>/dev/null | tail -1 | cut -d= -f2 || echo "")
        INIT_MS=$(grep "INIT_MS=" "$LOG.times" 2>/dev/null | tail -1 | cut -d= -f2 || echo "")
        K8S_MS=$(grep "K8S_MS=" "$LOG.times" 2>/dev/null | tail -1 | cut -d= -f2 || echo "")

        [ -n "$KERNEL_MS" ] && BOOT_TIMES+=("$KERNEL_MS")
        [ -n "$INIT_MS" ] && INIT_TIMES+=("$INIT_MS")
        [ -n "$K8S_MS" ] && K8S_TIMES+=("$K8S_MS")
    fi

    rm -f "$LOG" "$LOG.times"
    echo "" >&2
done

# Compute averages
avg() {
    local arr=("$@")
    if [ ${#arr[@]} -eq 0 ]; then
        echo "null"
        return
    fi
    local sum=0
    for v in "${arr[@]}"; do
        sum=$((sum + v))
    done
    echo $((sum / ${#arr[@]}))
}

# Image size
IMAGE_SIZE=$(stat -f%z "$IMAGE" 2>/dev/null || stat -c%s "$IMAGE" 2>/dev/null || echo 0)
IMAGE_SIZE_MB=$((IMAGE_SIZE / 1024 / 1024))

AVG_BOOT=$(avg "${BOOT_TIMES[@]+"${BOOT_TIMES[@]}"}")
AVG_INIT=$(avg "${INIT_TIMES[@]+"${INIT_TIMES[@]}"}")
AVG_K8S=$(avg "${K8S_TIMES[@]+"${K8S_TIMES[@]}"}")

echo "=== Results ===" >&2
echo "Image size:          ${IMAGE_SIZE_MB} MB" >&2
echo "Avg kernel load:     ${AVG_BOOT}ms" >&2
echo "Avg init complete:   ${AVG_INIT}ms" >&2
echo "Avg K8s Ready:       ${AVG_K8S}ms" >&2
echo "" >&2

# Output JSON
cat << EOF
{
  "benchmark": "kubesolo-os-boot",
  "image": "$(basename "$IMAGE")",
  "image_size_bytes": $IMAGE_SIZE,
  "image_size_mb": $IMAGE_SIZE_MB,
  "runs": $RUNS,
  "results": {
    "kernel_load_ms": $AVG_BOOT,
    "init_complete_ms": $AVG_INIT,
    "k8s_ready_ms": $AVG_K8S
  },
  "raw_kernel_ms": [$(IFS=,; echo "${BOOT_TIMES[*]+"${BOOT_TIMES[*]}"}")],
  "raw_init_ms": [$(IFS=,; echo "${INIT_TIMES[*]+"${INIT_TIMES[*]}"}")],
  "raw_k8s_ms": [$(IFS=,; echo "${K8S_TIMES[*]+"${K8S_TIMES[*]}"}")],
  "qemu_config": {
    "memory_mb": 1024,
    "cpus": 2,
    "kvm": $([ -e /dev/kvm ] && echo "true" || echo "false")
  }
}
EOF