#!/usr/bin/env bash
# test-boot-reliability - Stress test IncusOS first-boot reliability
#
# Tests whether IncusOS nodes reliably come up after installation by deploying
# multiple single-node VMs and checking port 8443 health.
#
# Three test groups with different seed configurations:
#   A) No update.yaml in seed (old behavior, state defaults only)
#   B) update.yaml with check_frequency: "6h" (current default)
#   C) update.yaml with check_frequency: "1h" (verifies seed takes effect)
#
# Runs batches of VMs in parallel, checks health, cleans up between batches.
# Reports per-group pass/fail rates at the end.
#
# Usage:
#   ./test-boot-reliability                    # 20 deploys in batches of 4
#   ./test-boot-reliability --count 12         # 12 deploys
#   ./test-boot-reliability --batch-size 2     # 2 parallel per batch
#   ./test-boot-reliability --dry-run          # Preview without deploying

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SEED_TOOL="${SCRIPT_DIR}/incusos-seed"

# Defaults
TOTAL_COUNT=20
BATCH_SIZE=4
START_VMID=850
DRY_RUN=false
HEALTH_TIMEOUT=150  # seconds to wait for port 8443 after boot
INSTALL_TIMEOUT=600 # seconds to wait for install to complete

# Proxmox connection (from proxmox.yaml)
PVE_HOST=""
PVE_NODE=""
PVE_STORAGE=""
PVE_ISO_STORAGE=""
PVE_BRIDGE=""
PVE_POOL=""
PVE_API_TOKEN_ID=""

# IncusOS ISO (must already be on Proxmox)
ISO_FILENAME=""

# Results directory (each background job writes a file)
RESULTS_DIR=""


# ---------------------------------------------------------------------------
# Colors
# ---------------------------------------------------------------------------

setup_colors() {
    if [[ "${NO_COLOR:-}" == "1" ]] || [[ "${TERM:-}" == "dumb" ]]; then
        BOLD="" RESET="" RED="" GREEN="" YELLOW="" CYAN="" DIM=""
    else
        BOLD="\033[1m" RESET="\033[0m" RED="\033[31m" GREEN="\033[32m"
        YELLOW="\033[33m" CYAN="\033[36m" DIM="\033[2m"
    fi
}
setup_colors

step()    { echo -e "${BOLD}[step]${RESET}    $*"; }
success() { echo -e "${GREEN}[ok]${RESET}      $*"; }
warn()    { echo -e "${YELLOW}[warn]${RESET}    $*"; }
error()   { echo -e "${RED}[error]${RESET}   $*"; }
info()    { echo -e "${DIM}[info]${RESET}    $*"; }
detail()  { echo -e "${DIM}          $*${RESET}"; }

# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

while [[ $# -gt 0 ]]; do
    case "$1" in
        --count)     TOTAL_COUNT="${2:?'--count requires a number'}"; shift 2 ;;
        --batch-size) BATCH_SIZE="${2:?'--batch-size requires a number'}"; shift 2 ;;
        --start-vmid) START_VMID="${2:?'--start-vmid requires a number'}"; shift 2 ;;
        --dry-run)   DRY_RUN=true; shift ;;
        -h|--help)
            echo "Usage: test-boot-reliability [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --count N        Total deploys (default: 20)"
            echo "  --batch-size N   Parallel VMs per batch (default: 4)"
            echo "  --start-vmid N   First VMID to use (default: 850)"
            echo "  --dry-run        Preview without deploying"
            echo "  -h, --help       Show this help"
            exit 0
            ;;
        *) error "Unknown option: $1"; exit 1 ;;
    esac
done

# ---------------------------------------------------------------------------
# Load environment
# ---------------------------------------------------------------------------

load_env() {
    [[ -n "${PROXMOX_TOKEN_SECRET:-}" ]] && return

    # Check envs/*/env in repo root
    local repo_root
    repo_root=$(cd "$SCRIPT_DIR/../.." && pwd)
    for env_file in "${repo_root}"/envs/*/env; do
        if [[ -f "$env_file" ]]; then
            # shellcheck disable=SC1091
            source "$env_file"
            [[ -n "${PROXMOX_TOKEN_SECRET:-}" ]] && return
        fi
    done

    # Walk up from script dir (backward compat)
    local dir="$SCRIPT_DIR"
    while [[ "$dir" != "/" ]]; do
        if [[ -f "${dir}/env" ]]; then
            # shellcheck disable=SC1091
            source "${dir}/env"
            return
        fi
        dir="$(dirname "$dir")"
    done
}
load_env

if [[ -z "${PROXMOX_TOKEN_SECRET:-}" ]]; then
    error "PROXMOX_TOKEN_SECRET not set. Create an env file or export it."
    exit 1
fi

# Load proxmox.yaml (simple key: value format, no nested structures)
yaml_get() {
    local file="$1" key="$2" default="${3:-}"
    local val
    val=$(awk -F': ' -v k="$key" '$1 == k {print $2; exit}' "$file" 2>/dev/null) || val=""
    # Strip surrounding quotes
    val="${val#\"}" ; val="${val%\"}"
    val="${val#\'}" ; val="${val%\'}"
    echo "${val:-$default}"
}

load_proxmox_yaml() {
    local yaml_file=""
    if [[ -f "${SCRIPT_DIR}/proxmox.yaml" ]]; then
        yaml_file="${SCRIPT_DIR}/proxmox.yaml"
    elif [[ -f "proxmox.yaml" ]]; then
        yaml_file="proxmox.yaml"
    else
        error "proxmox.yaml not found"
        exit 1
    fi

    PVE_HOST=$(yaml_get "$yaml_file" "host")
    PVE_NODE=$(yaml_get "$yaml_file" "node" "pve")
    PVE_STORAGE=$(yaml_get "$yaml_file" "storage" "local-lvm")
    PVE_ISO_STORAGE=$(yaml_get "$yaml_file" "iso_storage" "local")
    PVE_BRIDGE=$(yaml_get "$yaml_file" "bridge" "vmbr0")
    PVE_POOL=$(yaml_get "$yaml_file" "pool")
    PVE_API_TOKEN_ID=$(yaml_get "$yaml_file" "api_token_id")

    detail "Host: ${PVE_HOST}, Node: ${PVE_NODE}, Storage: ${PVE_STORAGE}"
    detail "Pool: ${PVE_POOL}, Token: ${PVE_API_TOKEN_ID}"
}

# ---------------------------------------------------------------------------
# Proxmox API helpers
# ---------------------------------------------------------------------------

pve_api() {
    local method="$1"
    local endpoint="$2"
    local data="${3:-}"
    local content_type="${4:-form}"

    local url="https://${PVE_HOST}:8006/api2/json${endpoint}"
    local auth="Authorization: PVEAPIToken=${PVE_API_TOKEN_ID}=${PROXMOX_TOKEN_SECRET}"

    local curl_args=(-fsSk -H "$auth")
    if [[ "$content_type" == "json" ]]; then
        curl_args+=(-H "Content-Type: application/json")
    fi

    case "$method" in
        GET)  curl "${curl_args[@]}" "$url" ;;
        POST) curl "${curl_args[@]}" -X POST ${data:+-d "$data"} "$url" ;;
        PUT)  curl "${curl_args[@]}" -X PUT ${data:+-d "$data"} "$url" ;;
        DELETE) curl "${curl_args[@]}" -X DELETE "$url" ;;
    esac
}

pve_vm_exists() {
    local vmid="$1"
    pve_api GET "/nodes/${PVE_NODE}/qemu/${vmid}/status/current" &>/dev/null
}

pve_vm_status() {
    local vmid="$1"
    pve_api GET "/nodes/${PVE_NODE}/qemu/${vmid}/status/current" 2>/dev/null | \
        python3 -c "import json,sys; print(json.load(sys.stdin)['data']['status'])" 2>/dev/null
}

pve_start_vm() {
    local vmid="$1"
    pve_api POST "/nodes/${PVE_NODE}/qemu/${vmid}/status/start" &>/dev/null
}

pve_stop_vm() {
    local vmid="$1"
    pve_api POST "/nodes/${PVE_NODE}/qemu/${vmid}/status/stop" &>/dev/null
}

pve_destroy_vm() {
    local vmid="$1"
    pve_api DELETE "/nodes/${PVE_NODE}/qemu/${vmid}?purge=1&destroy-unreferenced-disks=1" &>/dev/null
}

pve_vm_disk_writes() {
    local vmid="$1"
    pve_api GET "/nodes/${PVE_NODE}/qemu/${vmid}/status/current" 2>/dev/null | \
        python3 -c "
import json,sys
d = json.load(sys.stdin).get('data',{})
bs = d.get('blockstat',{}).get('scsi0',{})
print(bs.get('wr_bytes',0))
" 2>/dev/null || echo "0"
}

pve_vm_mac() {
    local vmid="$1"
    pve_api GET "/nodes/${PVE_NODE}/qemu/${vmid}/config" 2>/dev/null | \
        python3 -c "
import json, sys, re
data = json.load(sys.stdin).get('data', {})
net0 = data.get('net0', '')
m = re.search(r'([0-9A-Fa-f:]{17})', net0)
print(m.group(1).lower() if m else '')
" 2>/dev/null
}

pve_vm_ip() {
    # ARP-based IP detection matching incusos-proxmox pve_vm_ip().
    # Strategy: flush stale ARP → broadcast ping → check ARP → subnet sweep fallback.
    local vmid="$1"
    local mac
    mac=$(pve_vm_mac "$vmid") || return 1
    [[ -z "$mac" ]] && return 1

    ip neigh flush nud stale 2>/dev/null || true

    local bridge_subnet
    bridge_subnet=$(ip -4 addr show | awk '/inet / && !/127\.0\.0/ {print $2; exit}')

    # Try broadcast ping first (fast, but many hosts ignore it)
    if [[ -n "$bridge_subnet" ]]; then
        local broadcast
        broadcast=$(python3 -c "import ipaddress; print(ipaddress.ip_interface('${bridge_subnet}').network.broadcast_address)" 2>/dev/null) || true
        if [[ -n "$broadcast" ]]; then
            ping -c 2 -W 1 -b "$broadcast" &>/dev/null || true
            sleep 1
        fi
    fi

    # Check ARP table (filter by state)
    local ip
    ip=$(ip neigh show | awk -v mac="$mac" 'tolower($0) ~ tolower(mac) && (/REACHABLE/ || /DELAY/ || /STALE/) {print $1; exit}') || true

    if [[ -n "$ip" ]] && ping -c 1 -W 1 "$ip" &>/dev/null; then
        echo "$ip"
        return 0
    fi

    # Broadcast didn't work — parallel ping sweep of the subnet
    if [[ -n "$bridge_subnet" ]]; then
        local hosts
        hosts=$(python3 -c "
import ipaddress
net = ipaddress.ip_interface('${bridge_subnet}').network
for h in net.hosts():
    print(h)
" 2>/dev/null)
        if [[ -n "$hosts" ]]; then
            while IFS= read -r host_ip; do
                ping -c 1 -W 1 "$host_ip" &>/dev/null &
            done <<< "$hosts"
            wait
            sleep 1

            ip=$(ip neigh show | awk -v mac="$mac" 'tolower($0) ~ tolower(mac) && (/REACHABLE/ || /DELAY/) {print $1; exit}') || true

            if [[ -n "$ip" ]] && ping -c 1 -W 1 "$ip" &>/dev/null; then
                echo "$ip"
                return 0
            fi
        fi
    fi

    return 1
}

# ---------------------------------------------------------------------------
# Find IncusOS ISO already on Proxmox
# ---------------------------------------------------------------------------

find_incusos_iso() {
    step "Finding IncusOS ISO on Proxmox"
    local iso_list
    iso_list=$(pve_api GET "/nodes/${PVE_NODE}/storage/${PVE_ISO_STORAGE}/content?content=iso" 2>/dev/null) || iso_list=""

    ISO_FILENAME=$(echo "$iso_list" | python3 -c "
import json, sys
data = json.load(sys.stdin).get('data', [])
isos = [d['volid'].split('/')[-1] for d in data if 'IncusOS_' in d.get('volid', '')]
if isos:
    isos.sort()
    print(isos[-1])
" 2>/dev/null) || ISO_FILENAME=""

    if [[ -z "$ISO_FILENAME" ]]; then
        error "No IncusOS ISO found on Proxmox storage '${PVE_ISO_STORAGE}'"
        error "Deploy a lab first to upload the ISO, or use incusos-proxmox --phase upload"
        exit 1
    fi

    success "Using ISO: ${ISO_FILENAME}"
}

# ---------------------------------------------------------------------------
# Seed generation per group
# ---------------------------------------------------------------------------

generate_seed() {
    local group="$1"  # A, B, or C
    local name="$2"
    local output="$3"

    case "$group" in
        A)
            # No update.yaml -- uses incusos-seed to generate correct format,
            # then removes update.yaml from the seed directory before ISO creation.
            # This tests whether omitting update.yaml causes boot failures.
            local tmpdir
            tmpdir=$(mktemp -d)

            # Generate full seed via incusos-seed to a tar, extract, remove update.yaml
            "$SEED_TOOL" --format tar --app incus --output "${tmpdir}/seed.tar" \
                --force-install --force-reboot --hostname "$name" \
                --defaults --quiet &>/dev/null

            local seeddir="${tmpdir}/seedfiles"
            mkdir -p "$seeddir"
            tar xf "${tmpdir}/seed.tar" -C "$seeddir"
            rm -f "${seeddir}/update.yaml"

            # Generate ISO from the modified seed directory
            if command -v genisoimage &>/dev/null; then
                genisoimage -V SEED_DATA -J -r -o "$output" "$seeddir" 2>/dev/null
            elif command -v mkisofs &>/dev/null; then
                mkisofs -V SEED_DATA -J -r -o "$output" "$seeddir" 2>/dev/null
            fi
            rm -rf "$tmpdir"
            ;;
        B)
            # update.yaml with check_frequency: "6h" (default/current behavior)
            "$SEED_TOOL" --format iso --app incus --output "$output" \
                --force-install --force-reboot --hostname "$name" \
                --defaults --quiet &>/dev/null
            ;;
        C)
            # update.yaml with check_frequency: "1h" (verify seed takes effect)
            # Uses incusos-seed to generate correct format, then patches
            # update.yaml to use a different check_frequency.
            local tmpdir
            tmpdir=$(mktemp -d)

            # Generate full seed via incusos-seed
            "$SEED_TOOL" --format tar --app incus --output "${tmpdir}/seed.tar" \
                --force-install --force-reboot --hostname "$name" \
                --defaults --quiet &>/dev/null

            local seeddir="${tmpdir}/seedfiles"
            mkdir -p "$seeddir"
            tar xf "${tmpdir}/seed.tar" -C "$seeddir"

            # Patch check_frequency to 1h
            cat > "${seeddir}/update.yaml" <<'YAML'
# Update configuration
# Patched: check_frequency set to 1h for test group C
channel: "stable"
auto_reboot: false
check_frequency: "1h"
YAML

            # Generate ISO from the modified seed directory
            if command -v genisoimage &>/dev/null; then
                genisoimage -V SEED_DATA -J -r -o "$output" "$seeddir" 2>/dev/null
            elif command -v mkisofs &>/dev/null; then
                mkisofs -V SEED_DATA -J -r -o "$output" "$seeddir" 2>/dev/null
            fi
            rm -rf "$tmpdir"
            ;;
    esac
}

# ---------------------------------------------------------------------------
# Upload seed ISO to Proxmox
# ---------------------------------------------------------------------------

upload_seed_iso() {
    local local_path="$1"
    local filename="$2"

    pve_api POST "/nodes/${PVE_NODE}/storage/${PVE_ISO_STORAGE}/upload" \
        "" >/dev/null 2>&1 || true

    # Upload via curl multipart
    local url="https://${PVE_HOST}:8006/api2/json/nodes/${PVE_NODE}/storage/${PVE_ISO_STORAGE}/upload"
    local auth="Authorization: PVEAPIToken=${PVE_API_TOKEN_ID}=${PROXMOX_TOKEN_SECRET}"

    curl -fsSk -H "$auth" \
        -F "content=iso" \
        -F "filename=@${local_path};filename=${filename}" \
        "$url" &>/dev/null
}

# ---------------------------------------------------------------------------
# Create VM (minimal IncusOS settings)
# ---------------------------------------------------------------------------

create_vm() {
    local vmid="$1"
    local name="$2"
    local seed_iso="$3"

    # Proxmox API requires JSON for VM creation -- comma-separated disk/net
    # specs are misinterpreted when sent as form-encoded data.
    local json_payload
    json_payload=$(python3 -c "
import json
payload = {
    'vmid': ${vmid},
    'name': '${name}',
    'description': '[incusos-lab:managed] test-boot-reliability',
    'ostype': 'l26',
    'bios': 'ovmf',
    'machine': 'q35',
    'efidisk0': '${PVE_STORAGE}:0,efitype=4m,pre-enrolled-keys=0',
    'tpmstate0': '${PVE_STORAGE}:1,version=v2.0',
    'cpu': 'host',
    'cores': 4,
    'memory': 4096,
    'balloon': 0,
    'scsihw': 'virtio-scsi-pci',
    'scsi0': '${PVE_STORAGE}:50,iothread=1',
    'ide2': '${PVE_ISO_STORAGE}:iso/${ISO_FILENAME},media=cdrom',
    'ide3': '${PVE_ISO_STORAGE}:iso/${seed_iso},media=cdrom',
    'net0': 'virtio,bridge=${PVE_BRIDGE}',
    'boot': 'order=ide2;scsi0',
    'agent': 1,
}
pool = '${PVE_POOL}'
if pool:
    payload['pool'] = pool
print(json.dumps(payload))
")

    local url="https://${PVE_HOST}:8006/api2/json/nodes/${PVE_NODE}/qemu"
    local auth="Authorization: PVEAPIToken=${PVE_API_TOKEN_ID}=${PROXMOX_TOKEN_SECRET}"

    local create_response
    create_response=$(curl -fsSk -H "$auth" -H "Content-Type: application/json" \
        -X POST -d "$json_payload" "$url" 2>/dev/null) || return 1

    # VM creation is async -- wait for the UPID task to complete
    local upid
    upid=$(echo "$create_response" | python3 -c "import json,sys; print(json.load(sys.stdin).get('data',''))" 2>/dev/null) || upid=""
    if [[ -n "$upid" ]]; then
        local encoded_upid
        encoded_upid=$(python3 -c "import urllib.parse; print(urllib.parse.quote('${upid}', safe=''))" 2>/dev/null)
        local wait_elapsed=0
        while [[ $wait_elapsed -lt 60 ]]; do
            sleep 2
            wait_elapsed=$((wait_elapsed + 2))
            local task_status
            task_status=$(pve_api GET "/nodes/${PVE_NODE}/tasks/${encoded_upid}/status" 2>/dev/null | \
                python3 -c "import json,sys; print(json.load(sys.stdin).get('data',{}).get('status',''))" 2>/dev/null) || task_status=""
            if [[ "$task_status" == "stopped" ]]; then
                break
            fi
        done
    fi
}

# ---------------------------------------------------------------------------
# Single VM deploy + health check
# ---------------------------------------------------------------------------

write_result() {
    local vmid="$1"
    local group="$2"
    local outcome="$3"
    local detail="$4"
    echo "${outcome}|${detail}" > "${RESULTS_DIR}/${vmid}-${group}.result"
}

deploy_and_check() {
    local vmid="$1"
    local name="$2"
    local group="$3"
    local seed_iso="$4"

    info "[$name] Group ${group}: creating VM (VMID ${vmid})"

    if [[ "$DRY_RUN" == true ]]; then
        info "[$name] [dry-run] Would create, install, check health"
        write_result "$vmid" "$group" "dry-run" "skipped"
        return
    fi

    # Create VM
    if ! create_vm "$vmid" "$name" "$seed_iso"; then
        error "[$name] Failed to create VM"
        write_result "$vmid" "$group" "ERROR" "create failed"
        return
    fi

    # Start VM
    if ! pve_start_vm "$vmid"; then
        error "[$name] Failed to start VM"
        write_result "$vmid" "$group" "ERROR" "start failed"
        return
    fi

    # Monitor installation (blockstat polling)
    local elapsed=0
    local prev_wr=0
    local writes_started=false
    local stable_count=0
    local install_done=false

    while [[ $elapsed -lt $INSTALL_TIMEOUT ]]; do
        sleep 5
        elapsed=$((elapsed + 5))

        local status
        status=$(pve_vm_status "$vmid" 2>/dev/null) || status="unknown"

        if [[ "$status" == "stopped" ]]; then
            install_done=true
            break
        fi

        if [[ "$status" == "running" ]]; then
            local wr
            wr=$(pve_vm_disk_writes "$vmid" 2>/dev/null) || wr=0
            if [[ $wr -gt $prev_wr ]]; then
                writes_started=true
                stable_count=0
            elif [[ "$writes_started" == true ]]; then
                stable_count=$((stable_count + 1))
                if [[ $stable_count -ge 3 ]]; then
                    install_done=true
                    break
                fi
            fi
            prev_wr=$wr
        fi
    done

    if [[ "$install_done" != true ]]; then
        warn "[$name] Install timeout after ${INSTALL_TIMEOUT}s"
        write_result "$vmid" "$group" "FAIL" "install timeout"
        return
    fi

    # Stop VM if still running
    local cur_status
    cur_status=$(pve_vm_status "$vmid" 2>/dev/null) || cur_status="unknown"
    if [[ "$cur_status" == "running" ]]; then
        pve_stop_vm "$vmid" || true
        sleep 5
    fi

    # Remove install media, set boot order (must use JSON to avoid
    # form-encoding ambiguity with commas and equals signs)
    pve_api PUT "/nodes/${PVE_NODE}/qemu/${vmid}/config" \
        '{"delete":"ide2,ide3","boot":"order=scsi0"}' json &>/dev/null || true

    # Start from disk
    if ! pve_start_vm "$vmid"; then
        error "[$name] Failed to start from disk"
        write_result "$vmid" "$group" "FAIL" "disk boot failed"
        return
    fi

    # Health check: wait for port 8443
    info "[$name] Waiting for port 8443..."
    local health_elapsed=0
    local healthy=false

    # Wait 45s for VM to boot and get DHCP lease.
    # Stagger by vmid mod to reduce ARP sweep contention across parallel jobs.
    local stagger=$(( (vmid % 4) * 5 ))
    sleep $((45 + stagger))
    health_elapsed=$((45 + stagger))

    # Detect IP with retries
    local ip=""
    local ip_retry=0
    while [[ $ip_retry -lt 6 ]] && [[ -z "$ip" ]]; do
        ip=$(pve_vm_ip "$vmid" 2>/dev/null) || ip=""
        if [[ -n "$ip" ]]; then
            break
        fi
        sleep 10
        health_elapsed=$((health_elapsed + 10))
        ip_retry=$((ip_retry + 1))
    done

    if [[ -z "$ip" ]]; then
        warn "[$name] Could not detect IP"
        write_result "$vmid" "$group" "FAIL" "no IP detected"
        return
    fi

    # Poll port 8443
    while [[ $health_elapsed -lt $HEALTH_TIMEOUT ]]; do
        if curl -sk --connect-timeout 3 "https://${ip}:8443" &>/dev/null; then
            healthy=true
            break
        fi
        sleep 5
        health_elapsed=$((health_elapsed + 5))
    done

    if [[ "$healthy" == true ]]; then
        success "[$name] Group ${group}: PASS — port 8443 up at ${ip} (${health_elapsed}s)"
        local result_detail="healthy at ${ip} in ${health_elapsed}s"

        # For group C, verify check_frequency via IncusOS API
        if [[ "$group" == "C" ]]; then
            local freq="?"
            local cert_file="${HOME}/.config/incus/client.crt"
            local key_file="${HOME}/.config/incus/client.key"
            if [[ -f "$cert_file" ]] && [[ -f "$key_file" ]]; then
                freq=$(curl -sk --cert "$cert_file" --key "$key_file" \
                    "https://${ip}:8443/1.0/system/update" 2>/dev/null | \
                    python3 -c "import json,sys; d=json.load(sys.stdin); print(d.get('metadata',d).get('check_frequency','?'))" 2>/dev/null) || freq="?"
            fi
            detail "[$name] check_frequency from API: ${freq}"
            result_detail="healthy at ${ip}, check_frequency=${freq}"
        fi
        write_result "$vmid" "$group" "PASS" "$result_detail"
    else
        warn "[$name] Group ${group}: FAIL — port 8443 not reachable at ${ip} after ${HEALTH_TIMEOUT}s"
        write_result "$vmid" "$group" "FAIL" "port 8443 unreachable at ${ip}"
    fi
}

# ---------------------------------------------------------------------------
# Cleanup a VM
# ---------------------------------------------------------------------------

cleanup_vm() {
    local vmid="$1"
    local name="$2"
    local seed_iso="$3"

    if [[ "$DRY_RUN" == true ]]; then
        return
    fi

    # Stop if running
    local status
    status=$(pve_vm_status "$vmid" 2>/dev/null) || status="unknown"
    if [[ "$status" == "running" ]]; then
        pve_stop_vm "$vmid" || true
        sleep 3
    fi

    # Destroy
    pve_destroy_vm "$vmid" || true

    # Remove seed ISO from storage
    pve_api DELETE "/nodes/${PVE_NODE}/storage/${PVE_ISO_STORAGE}/content/${PVE_ISO_STORAGE}:iso/${seed_iso}" &>/dev/null || true

    # Remove incus remote
    if command -v incus &>/dev/null; then
        local current_remote
        current_remote=$(incus remote get-default 2>/dev/null) || true
        if [[ "$current_remote" == "$name" ]]; then
            incus remote switch local 2>/dev/null || true
        fi
        incus remote remove "$name" 2>/dev/null || true
    fi
}

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

echo -e "${BOLD}IncusOS Boot Reliability Test${RESET}"
echo ""
echo "Test design:"
echo "  Group A: no update.yaml in seed (state defaults only)"
echo "  Group B: update.yaml with check_frequency: \"6h\" (current default)"
echo "  Group C: update.yaml with check_frequency: \"1h\" (verify seed effect)"
echo ""
echo "Deploys: ${TOTAL_COUNT} total, ${BATCH_SIZE} parallel per batch"
echo "VMIDs: ${START_VMID}-$((START_VMID + TOTAL_COUNT - 1))"
echo ""

step "Loading configuration"
load_proxmox_yaml
find_incusos_iso

# Assign groups: distribute evenly across A, B, C
# Round-robin: 0→A, 1→B, 2→C, 3→A, 4→B, ...
declare -a VM_GROUPS=()
declare -a VM_NAMES=()
declare -a VM_VMIDS=()
declare -a VM_SEEDS=()

groups=("A" "B" "C")
for ((idx=0; idx<TOTAL_COUNT; idx++)); do
    group="${groups[$((idx % 3))]}"
    vmid=$((START_VMID + idx))
    name="boot-test-${group,,}-${vmid}"
    seed_name="seed-${name}.iso"

    VM_GROUPS+=("$group")
    VM_NAMES+=("$name")
    VM_VMIDS+=("$vmid")
    VM_SEEDS+=("$seed_name")
done

echo ""
echo -e "${BOLD}Test matrix:${RESET}"
for ((idx=0; idx<TOTAL_COUNT; idx++)); do
    echo "  VMID ${VM_VMIDS[$idx]}: ${VM_NAMES[$idx]} (Group ${VM_GROUPS[$idx]})"
done

echo ""

# Check for VMID collisions
step "Checking for VMID collisions"
if [[ "$DRY_RUN" != true ]]; then
    for ((idx=0; idx<TOTAL_COUNT; idx++)); do
        if pve_vm_exists "${VM_VMIDS[$idx]}"; then
            error "VMID ${VM_VMIDS[$idx]} already exists. Choose a different --start-vmid range."
            exit 1
        fi
    done
    success "No collisions"
fi

# Generate and upload seed ISOs
step "Generating seed ISOs"
SEED_TMPDIR=$(mktemp -d)
RESULTS_DIR=$(mktemp -d)
trap 'rm -rf "$SEED_TMPDIR" "$RESULTS_DIR"' EXIT

for ((idx=0; idx<TOTAL_COUNT; idx++)); do
    local_seed="${SEED_TMPDIR}/${VM_SEEDS[$idx]}"
    generate_seed "${VM_GROUPS[$idx]}" "${VM_NAMES[$idx]}" "$local_seed"

    if [[ "$DRY_RUN" == true ]]; then
        info "[dry-run] Would generate seed: ${VM_SEEDS[$idx]} (Group ${VM_GROUPS[$idx]})"
    else
        upload_seed_iso "$local_seed" "${VM_SEEDS[$idx]}"
        detail "Uploaded: ${VM_SEEDS[$idx]} (Group ${VM_GROUPS[$idx]})"
    fi
done
success "All seed ISOs generated and uploaded"

# Run batches
echo ""
total_batches=$(( (TOTAL_COUNT + BATCH_SIZE - 1) / BATCH_SIZE ))
batch_num=0

for ((batch_start=0; batch_start<TOTAL_COUNT; batch_start+=BATCH_SIZE)); do
    batch_num=$((batch_num + 1))
    batch_end=$((batch_start + BATCH_SIZE))
    [[ $batch_end -gt $TOTAL_COUNT ]] && batch_end=$TOTAL_COUNT
    batch_count=$((batch_end - batch_start))

    echo ""
    step "Batch ${batch_num}/${total_batches}: deploying ${batch_count} VMs (VMIDs $((START_VMID + batch_start))-$((START_VMID + batch_end - 1)))"
    echo ""

    # Deploy all VMs in this batch in parallel
    pids=()
    for ((idx=batch_start; idx<batch_end; idx++)); do
        deploy_and_check "${VM_VMIDS[$idx]}" "${VM_NAMES[$idx]}" \
            "${VM_GROUPS[$idx]}" "${VM_SEEDS[$idx]}" &
        pids+=($!)
    done

    # Wait for all in this batch to complete
    for pid in "${pids[@]}"; do
        wait "$pid" || true
    done

    echo ""
    step "Batch ${batch_num} complete — cleaning up"

    # Cleanup all VMs in this batch
    for ((idx=batch_start; idx<batch_end; idx++)); do
        cleanup_vm "${VM_VMIDS[$idx]}" "${VM_NAMES[$idx]}" "${VM_SEEDS[$idx]}"
        detail "Cleaned: ${VM_NAMES[$idx]}"
    done
    success "Batch ${batch_num} cleaned"
done

# ---------------------------------------------------------------------------
# Results summary (read from result files)
# ---------------------------------------------------------------------------

echo ""
echo ""
echo -e "${BOLD}========================================${RESET}"
echo -e "${BOLD}  Test Results Summary${RESET}"
echo -e "${BOLD}========================================${RESET}"
echo ""

# Per-group stats
for group in A B C; do
    total=0
    pass=0
    fail=0
    err=0
    for ((idx=0; idx<TOTAL_COUNT; idx++)); do
        if [[ "${VM_GROUPS[$idx]}" == "$group" ]]; then
            total=$((total + 1))
            rfile="${RESULTS_DIR}/${VM_VMIDS[$idx]}-${group}.result"
            if [[ -f "$rfile" ]]; then
                outcome=$(cut -d'|' -f1 "$rfile")
                case "$outcome" in
                    PASS)    pass=$((pass + 1)) ;;
                    FAIL)    fail=$((fail + 1)) ;;
                    ERROR)   err=$((err + 1)) ;;
                esac
            else
                err=$((err + 1))  # no result file = something went wrong
            fi
        fi
    done

    label=""
    case "$group" in
        A) label="No update.yaml (state defaults)" ;;
        B) label="update.yaml check_frequency=6h" ;;
        C) label="update.yaml check_frequency=1h" ;;
    esac

    if [[ $fail -eq 0 ]] && [[ $err -eq 0 ]]; then
        echo -e "  Group ${group} (${label}):"
        echo -e "    ${GREEN}${pass}/${total} PASS${RESET}"
    else
        echo -e "  Group ${group} (${label}):"
        echo -e "    ${GREEN}${pass} PASS${RESET}, ${RED}${fail} FAIL${RESET}, ${RED}${err} ERROR${RESET} (of ${total})"
    fi
done

echo ""
echo -e "${BOLD}Per-VM details:${RESET}"
for ((idx=0; idx<TOTAL_COUNT; idx++)); do
    rfile="${RESULTS_DIR}/${VM_VMIDS[$idx]}-${VM_GROUPS[$idx]}.result"
    outcome="MISSING"
    rdetail="no result file"
    if [[ -f "$rfile" ]]; then
        outcome=$(cut -d'|' -f1 "$rfile")
        rdetail=$(cut -d'|' -f2- "$rfile")
    fi

    marker=""
    case "$outcome" in
        PASS)    marker="${GREEN}PASS${RESET}" ;;
        FAIL)    marker="${RED}FAIL${RESET}" ;;
        ERROR)   marker="${RED}ERROR${RESET}" ;;
        *)       marker="${DIM}${outcome}${RESET}" ;;
    esac
    echo -e "  ${VM_VMIDS[$idx]} ${VM_NAMES[$idx]} Group ${VM_GROUPS[$idx]}: ${marker} -- ${rdetail}"
done

echo ""

# Overall assessment
total_fail=0
for ((idx=0; idx<TOTAL_COUNT; idx++)); do
    rfile="${RESULTS_DIR}/${VM_VMIDS[$idx]}-${VM_GROUPS[$idx]}.result"
    if [[ -f "$rfile" ]]; then
        outcome=$(cut -d'|' -f1 "$rfile")
        if [[ "$outcome" == "FAIL" ]] || [[ "$outcome" == "ERROR" ]]; then
            total_fail=$((total_fail + 1))
        fi
    else
        total_fail=$((total_fail + 1))
    fi
done

if [[ $total_fail -eq 0 ]]; then
    success "All ${TOTAL_COUNT} deploys passed — no boot failures detected"
else
    warn "${total_fail}/${TOTAL_COUNT} deploys failed"
    echo ""
    echo "Analysis:"
    for group in A B C; do
        g_fail=0
        g_total=0
        for ((idx=0; idx<TOTAL_COUNT; idx++)); do
            if [[ "${VM_GROUPS[$idx]}" == "$group" ]]; then
                g_total=$((g_total + 1))
                rfile="${RESULTS_DIR}/${VM_VMIDS[$idx]}-${group}.result"
                outcome=""
                [[ -f "$rfile" ]] && outcome=$(cut -d'|' -f1 "$rfile")
                if [[ "$outcome" == "FAIL" ]] || [[ "$outcome" == "ERROR" ]] || [[ -z "$outcome" ]]; then
                    g_fail=$((g_fail + 1))
                fi
            fi
        done
        echo "  Group ${group}: ${g_fail}/${g_total} failures"
    done
fi

echo ""
echo "Done."