#!/bin/bash
# SPDX-License-Identifier: GPL-3.0-or-later
# Run maintenance-console remedy scenarios in the test VM
# Author: Craig Jennings <craigmartinjennings@gmail.com>
# License: GNU GPLv3
#
# Layer 3 of the maint test strategy: each scenario breaks the VM in a known
# way over ssh, runs the real `maint fix <id>` inside the VM, and asserts the
# post-state. Scenarios live in maint-scenarios/*.sh and declare a GROUP;
# scenarios in one group are non-conflicting and share a VM boot, and a
# stop -> restore -> boot cycle runs only between groups (snapshot ops need
# the VM stopped). A scenario that fails may leave residue behind for its
# in-group successors -- the group restore boundary contains it, so read a
# multi-failure group front to back.
#
# Boot plan:
#   1. restore clean-install, boot, bootstrap (packages + maint tree + TOML)
#   2. stop, snapshot maint-ready
#   3. per group: boot, run scenarios, stop, restore maint-ready
#   4. restore clean-install, delete maint-ready (base stays pristine)
#
# Usage: run-maint-scenarios.sh [--list] [--group NAME] [--keep]
#   --list   print the validated batch plan and exit (no VM, no KVM needed)
#   --group  run (or list) only the named group
#   --keep   keep the VM running in its post-run state; skip the restores
# Env: FS_PROFILE=btrfs|zfs        base image + scenario profile filter
#      MAINT_SCENARIO_DIR          scenario dir (default: maint-scenarios/)
#      MAINT_SRC                   maint package source tree
#                                  (default: ~/.dotfiles/maint/src/maint)

set -e

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

source "$SCRIPT_DIR/lib/logging.sh"
source "$SCRIPT_DIR/lib/vm-utils.sh"

SCENARIO_DIR="${MAINT_SCENARIO_DIR:-$SCRIPT_DIR/maint-scenarios}"
MAINT_SRC="${MAINT_SRC:-$HOME/.dotfiles/maint/src/maint}"
THRESHOLDS_TOML="$PROJECT_ROOT/configs/maintenance-thresholds.toml"
ROOT_PASSWORD="archsetup"
READY_SNAPSHOT="maint-ready"

LIST_ONLY=false
ONLY_GROUP=""
KEEP_VM=false

usage() {
    echo "Usage: $0 [--list] [--group NAME] [--keep]"
    echo "  --list   print the validated scenario plan and exit (no VM)"
    echo "  --group  run only the named scenario group"
    echo "  --keep   keep the VM running after the run (skip restores)"
    echo "Env: FS_PROFILE=btrfs|zfs, MAINT_SCENARIO_DIR, MAINT_SRC"
}

while [[ $# -gt 0 ]]; do
    case $1 in
        --list)  LIST_ONLY=true; shift ;;
        --group) ONLY_GROUP="${2:?--group requires a value}"; shift 2 ;;
        --keep)  KEEP_VM=true; shift ;;
        *) usage; exit 1 ;;
    esac
done

# ─── Plan: enumerate, validate, group ────────────────────────────────
# The plan layer is pure -- no VM, no KVM -- so --list works anywhere and
# the unit suite (tests/maint-scenarios/) exercises it directly.

# Parallel arrays over scenarios, in file order.
S_FILES=() S_NAMES=() S_DESCS=() S_GROUPS=()
# Groups in first-appearance order.
GROUP_ORDER=()

# Extract one declared var from a scenario file by sourcing it in a clean
# subshell. Scenario files must only define vars + functions; any top-level
# command fails here (no helpers exist), which is the contract.
_scenario_var() {  # <file> <varname>
    bash -c 'set -eu; source "$1" || exit 9; eval "printf %s \"\${$2-}\""' \
        _probe "$1" "$2" 2>/dev/null
}

_validate_scenario() {  # <file> -> fatal on contract violation
    local f="$1" base probe_err
    base=$(basename "$f")
    probe_err=$(bash -c 'set -eu; source "$1"
        : "${SCENARIO_DESC:?missing SCENARIO_DESC}"
        : "${SCENARIO_GROUP:?missing SCENARIO_GROUP}"
        : "${SCENARIO_PROFILES:?missing SCENARIO_PROFILES}"
        for p in $SCENARIO_PROFILES; do
            case "$p" in btrfs|zfs|any) ;; *)
                echo "bad profile token: $p" >&2; exit 1 ;;
            esac
        done
        declare -f scenario_break scenario_fix scenario_assert >/dev/null \
            || { echo "missing scenario_break/fix/assert" >&2; exit 1; }' \
        _probe "$f" 2>&1 >/dev/null) \
        || fatal "scenario contract violation in $base: $probe_err"
}

_profile_matches() {  # <profiles> -> 0 if scenario runs under $FS_PROFILE
    local p
    for p in $1; do
        [ "$p" = "any" ] || [ "$p" = "$FS_PROFILE" ] && return 0
    done
    return 1
}

build_plan() {
    [ -d "$SCENARIO_DIR" ] || fatal "scenario dir not found: $SCENARIO_DIR"
    local f base name desc group profiles seen_groups=" "
    local any=false
    for f in "$SCENARIO_DIR"/*.sh; do
        [ -e "$f" ] || break
        any=true
        _validate_scenario "$f"
        base=$(basename "$f" .sh)
        name="${base#[0-9][0-9]-}"
        desc=$(_scenario_var "$f" SCENARIO_DESC)
        group=$(_scenario_var "$f" SCENARIO_GROUP)
        profiles=$(_scenario_var "$f" SCENARIO_PROFILES)
        _profile_matches "$profiles" || continue
        [ -n "$ONLY_GROUP" ] && [ "$group" != "$ONLY_GROUP" ] && continue
        S_FILES+=("$f"); S_NAMES+=("$name")
        S_DESCS+=("$desc"); S_GROUPS+=("$group")
        case "$seen_groups" in
            *" $group "*) ;;
            *) GROUP_ORDER+=("$group"); seen_groups="$seen_groups$group " ;;
        esac
    done
    [ "$any" = "true" ] || fatal "no scenario files in $SCENARIO_DIR"
    if [ -n "$ONLY_GROUP" ] && [ ${#S_FILES[@]} -eq 0 ]; then
        fatal "no scenarios match group '$ONLY_GROUP' (profile $FS_PROFILE)"
    fi
    [ ${#S_FILES[@]} -gt 0 ] \
        || fatal "no scenarios match profile $FS_PROFILE"
}

print_plan() {
    local g i
    echo "maint scenario plan (profile: $FS_PROFILE)"
    for g in "${GROUP_ORDER[@]}"; do
        echo "group $g:"
        for i in "${!S_FILES[@]}"; do
            [ "${S_GROUPS[$i]}" = "$g" ] \
                && echo "  ${S_NAMES[$i]} — ${S_DESCS[$i]}"
        done
    done
}

# --list needs no logging file, VM, or KVM. Give fatal() a LOGFILE target
# only when we actually run.
if [ "$LIST_ONLY" = "true" ]; then
    LOGFILE=/dev/null
    build_plan
    print_plan
    exit 0
fi

# ─── Run ─────────────────────────────────────────────────────────────

TIMESTAMP=$(date +'%Y%m%d-%H%M%S')
RESULTS_DIR="$PROJECT_ROOT/test-results/maint-$TIMESTAMP"
mkdir -p "$RESULTS_DIR"
LOGFILE="$RESULTS_DIR/scenarios.log"
init_logging "$LOGFILE"
init_vm_paths "$PROJECT_ROOT/vm-images"

build_plan

section "Maint Remedy Scenarios: $TIMESTAMP"
info "Profile: $FS_PROFILE (image: $(basename "$DISK_PATH"))"
info "Maint source: $MAINT_SRC"
print_plan | tee -a "$LOGFILE"

[ -f "$DISK_PATH" ] || fatal "base disk not found: $DISK_PATH — build it: FS_PROFILE=$FS_PROFILE $SCRIPT_DIR/create-base-vm.sh"
snapshot_exists "$DISK_PATH" "clean-install" \
    || fatal "snapshot 'clean-install' not found on $DISK_PATH"
[ -d "$MAINT_SRC" ] || fatal "maint source tree not found: $MAINT_SRC"
[ -f "$THRESHOLDS_TOML" ] || fatal "thresholds TOML not found: $THRESHOLDS_TOML"
check_prerequisites || fatal "missing prerequisites"

CLEANUP_DONE=0
cleanup_scenarios() {
    [ "$CLEANUP_DONE" = "1" ] && return 0
    CLEANUP_DONE=1
    [ "$KEEP_VM" = "true" ] && return 0
    # Never silent: a failed restore/delete here leaves the base image dirty
    # for the next `make test`, so name it even though the trap can't abort.
    stop_qemu 2>/dev/null \
        || echo "[!] cleanup: stop_qemu failed — VM may still hold the image" >&2
    restore_snapshot "$DISK_PATH" "clean-install" 2>/dev/null \
        || echo "[!] cleanup: clean-install restore FAILED — base image is dirty" >&2
    if snapshot_exists "$DISK_PATH" "$READY_SNAPSHOT"; then
        delete_snapshot "$DISK_PATH" "$READY_SNAPSHOT" 2>/dev/null \
            || echo "[!] cleanup: stray '$READY_SNAPSHOT' snapshot left on the base image" >&2
    fi
}
trap cleanup_scenarios EXIT

# ─── Target helpers (available to scenario scripts) ──────────────────

mexec() {  # run a command in the VM as root
    vm_exec "$ROOT_PASSWORD" "$@"
}

# maint invocation inside the VM. Root runs the verbs directly, so
# MAINT_SUDO is empty (same switch the unit suites use).
mmaint() {
    mexec "cd /root/maint-pkg && MAINT_SUDO= PYTHONPATH=/root/maint-pkg python3 -m maint $*"
}

# Run one remedy: dry-run first (argv parity gate — the printed argv is the
# arm-press string), then for real. Fails if either exits non-zero.
mfix() {
    local id="$1"; shift || true
    step "maint fix $id ${*:+$* }--dry-run"
    mmaint fix "$id" "$@" --dry-run >> "$LOGFILE" 2>&1 \
        || { error "dry-run failed: maint fix $id $*"; return 1; }
    step "maint fix $id $*"
    mmaint fix "$id" "$@" 2>&1 | tee -a "$LOGFILE"
    return "${PIPESTATUS[0]}"
}

# Assert a metric's severity in `maint status --json`.
massert_metric() {  # <metric-id> <expected-severity>
    local mid="$1" want="$2" got
    got=$(mmaint status --json | python3 -c "
import json, sys
env = json.load(sys.stdin)
ms = [m for m in env['metrics'] if m['id'] == '$mid']
print(ms[0]['severity'] if ms else 'MISSING')
")
    if [ "$got" = "$want" ]; then
        success "metric $mid is $want"
        return 0
    fi
    error "metric $mid: expected $want, got $got"
    return 1
}

# ─── Bootstrap ───────────────────────────────────────────────────────

boot_vm() {
    start_qemu "$DISK_PATH" "disk" "" "none" || fatal "failed to start VM"
    wait_for_ssh "$ROOT_PASSWORD" 180 || fatal "VM SSH not available"
}

bootstrap_vm() {
    section "Bootstrap: packages + maint tree + thresholds"
    step "Refreshing keyring + installing runtime deps"
    # The base image can be months old: refresh the keyring first so the
    # dep install doesn't die on expired signatures.
    mexec "pacman -Sy --noconfirm archlinux-keyring" >> "$LOGFILE" 2>&1 \
        || fatal "keyring refresh failed (network up?)"
    mexec "pacman -S --noconfirm python pacman-contrib expac cronie" \
        >> "$LOGFILE" 2>&1 || fatal "dep install failed"

    step "Copying maint package tree"
    local tarball
    tarball=$(mktemp)
    tar -C "$(dirname "$MAINT_SRC")" --exclude='__pycache__' -czf "$tarball" \
        "$(basename "$MAINT_SRC")"
    copy_to_vm "$tarball" "/tmp/maint-src.tgz" "$ROOT_PASSWORD"
    rm -f "$tarball"
    mexec "rm -rf /root/maint-pkg && mkdir -p /root/maint-pkg && tar -C /root/maint-pkg -xzf /tmp/maint-src.tgz && rm /tmp/maint-src.tgz" \
        >> "$LOGFILE" 2>&1 || fatal "maint tree extract failed"

    step "Installing thresholds TOML"
    copy_to_vm "$THRESHOLDS_TOML" "/tmp/maintenance-thresholds.toml" \
        "$ROOT_PASSWORD"
    mexec "mkdir -p /root/.config/archsetup && mv /tmp/maintenance-thresholds.toml /root/.config/archsetup/maintenance-thresholds.toml" \
        >> "$LOGFILE" 2>&1 || fatal "TOML install failed"

    step "Smoke: maint status --json runs"
    mmaint status --json > "$RESULTS_DIR/bootstrap-status.json" \
        || fatal "maint status failed in the VM"
    success "Bootstrap complete ($(python3 -c "
import json; print(len(json.load(open('$RESULTS_DIR/bootstrap-status.json'))['metrics']))") metrics probed)"
}

# ─── Execution ───────────────────────────────────────────────────────

PASS=() FAIL=()

run_scenario() {  # <index>
    local i="$1" name="${S_NAMES[$1]}" f="${S_FILES[$1]}"
    section "Scenario: $name — ${S_DESCS[$1]}"
    # Sourced in the runner's shell so break/fix/assert see the helpers;
    # unset afterwards so the next scenario can't inherit stale functions.
    # shellcheck disable=SC1090
    source "$f"
    local ok=true
    step "break"
    if ! scenario_break; then
        error "$name: break step failed"
        ok=false
    fi
    if [ "$ok" = "true" ]; then
        step "fix"
        if ! scenario_fix; then
            error "$name: fix step failed"
            ok=false
        fi
    fi
    if [ "$ok" = "true" ]; then
        step "assert"
        if ! scenario_assert; then
            error "$name: post-state assertion failed"
            ok=false
        fi
    fi
    unset -f scenario_break scenario_fix scenario_assert
    unset SCENARIO_DESC SCENARIO_GROUP SCENARIO_PROFILES
    if [ "$ok" = "true" ]; then
        success "PASS: $name"
        PASS+=("$name")
    else
        error "FAIL: $name"
        FAIL+=("$name")
    fi
}

section "Preparing VM"
stop_qemu 2>/dev/null || true
step "Restoring clean-install snapshot"
restore_snapshot "$DISK_PATH" "clean-install" || fatal "restore failed"
boot_vm
bootstrap_vm

step "Freezing bootstrap as snapshot: $READY_SNAPSHOT"
stop_qemu
snapshot_exists "$DISK_PATH" "$READY_SNAPSHOT" \
    && delete_snapshot "$DISK_PATH" "$READY_SNAPSHOT"
create_snapshot "$DISK_PATH" "$READY_SNAPSHOT" || fatal "snapshot failed"

first_group=true
for g in "${GROUP_ORDER[@]}"; do
    section "Group: $g"
    if [ "$first_group" = "true" ]; then
        first_group=false
    else
        step "Isolation boundary: restore $READY_SNAPSHOT"
        stop_qemu
        restore_snapshot "$DISK_PATH" "$READY_SNAPSHOT" || fatal "restore failed"
    fi
    boot_vm
    for i in "${!S_FILES[@]}"; do
        [ "${S_GROUPS[$i]}" = "$g" ] && run_scenario "$i"
    done
    stop_qemu
done

# ─── Report ──────────────────────────────────────────────────────────

section "Results"
for n in "${PASS[@]}"; do success "PASS $n"; done
for n in "${FAIL[@]}"; do error "FAIL $n"; done
info "Log: $LOGFILE"

if [ "$KEEP_VM" = "true" ]; then
    warn "--keep: VM left in post-run state; base image NOT restored"
    warn "restore by hand: qemu-img snapshot -a clean-install $DISK_PATH"
    CLEANUP_DONE=1
else
    cleanup_scenarios
fi

if [ ${#FAIL[@]} -gt 0 ]; then
    error "${#FAIL[@]} scenario(s) failed, ${#PASS[@]} passed"
    exit 1
fi
success "All ${#PASS[@]} scenarios passed"
exit 0