#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────
# Test: reprovision preserves /home, /srv, /var/lib/longhorn
#
# Usage:  sudo bash test-reprovision.sh
#         sudo bash test-reprovision.sh --skip-first-install  # if disk already has a first install
#         sudo bash test-reprovision.sh --cleanup              # just remove the VM and disk
# ─────────────────────────────────────────────────────────────
set -euo pipefail

VM_NAME="test-bastion-ks"
DISK_PATH="/var/lib/libvirt/images/test-reprovision.qcow2"
DISK_SIZE=20  # GB
KS_PATH="/tmp/test-vm.ks"
FEDORA_MIRROR="https://download.fedoraproject.org/pub/fedora/linux/releases/43/Everything/x86_64/os/"
OVMF_CODE="/usr/share/edk2/ovmf/OVMF_CODE.fd"
OVMF_VARS="/usr/share/OVMF/OVMF_VARS.fd"

RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'

log()  { echo -e "${GREEN}[test]${NC} $*"; }
err()  { echo -e "${RED}[test]${NC} $*" >&2; }
step() { echo -e "\n${CYAN}${BOLD}══ $* ══${NC}\n"; }

cleanup_vm() {
    virsh destroy "$VM_NAME" 2>/dev/null || true
    virsh undefine "$VM_NAME" --nvram 2>/dev/null || true
}

cleanup_all() {
    cleanup_vm
    rm -f "$DISK_PATH"
    log "Cleaned up VM and disk"
}

# ── Handle args ──
SKIP_FIRST=false
for arg in "$@"; do
    case "$arg" in
        --skip-first-install) SKIP_FIRST=true ;;
        --cleanup) cleanup_all; exit 0 ;;
    esac
done

[[ $EUID -eq 0 ]] || { err "Must run as root"; exit 1; }

# ── Generate kickstart ──
generate_kickstart() {
    cat > "$KS_PATH" << 'KSEOF'
text
reboot
lang en_GB.UTF-8
keyboard uk
timezone Europe/London --utc
network --bootproto=dhcp --activate --hostname=test-vm.ad.itaz.eu
rootpw --plaintext testpass
user --name=michal --groups=wheel
bootloader --append="console=ttyS0,115200n8"
url --mirrorlist=https://mirrors.fedoraproject.org/mirrorlist?repo=fedora-43&arch=x86_64
%include /tmp/part.ks

%pre --log=/tmp/pre-partition.log
#!/bin/bash
set -x
VG="labvg"
DISK="vda"

REPROVISION=no
if vgs $VG &>/dev/null; then
    echo "=== REPROVISION MODE ==="
    REPROVISION=yes
    PRESERVE_LONGHORN=no; PRESERVE_SRV=no; PRESERVE_HOME=no
    lvs $VG/longhorn &>/dev/null && PRESERVE_LONGHORN=yes
    lvs $VG/srv      &>/dev/null && PRESERVE_SRV=yes
    lvs $VG/home     &>/dev/null && PRESERVE_HOME=yes
    echo "Preserving: longhorn=$PRESERVE_LONGHORN srv=$PRESERVE_SRV home=$PRESERVE_HOME"
    for lv in root var varlog swap; do
        lvremove -f $VG/$lv 2>/dev/null || true
    done
fi

if [ "$REPROVISION" = "yes" ]; then
    EFI_PART=$(blkid -t TYPE=vfat -o device /dev/${DISK}* 2>/dev/null | head -1)
    BOOT_PART=$(blkid -t TYPE=ext4 -o device /dev/${DISK}* 2>/dev/null | head -1)
    EFI_PART=${EFI_PART:-/dev/${DISK}1}
    BOOT_PART=${BOOT_PART:-/dev/${DISK}2}
    echo "Reusing EFI=$EFI_PART BOOT=$BOOT_PART"

    cat > /tmp/part.ks << PARTEOF
ignoredisk --only-use=$DISK
clearpart --none
part /boot/efi --onpart=$EFI_PART --fstype=efi
part /boot --onpart=$BOOT_PART --fstype=ext4
volgroup labvg --useexisting --noformat
logvol swap --vgname=labvg --name=swap --fstype=swap --size=1024
logvol / --vgname=labvg --name=root --fstype=xfs --size=4096
logvol /var --vgname=labvg --name=var --fstype=xfs --size=3072
logvol /var/log --vgname=labvg --name=varlog --fstype=xfs --size=1024
PARTEOF
    if [ "$PRESERVE_HOME" = "yes" ]; then
        echo "logvol /home --vgname=labvg --name=home --useexisting --noformat" >> /tmp/part.ks
    else
        echo "logvol /home --vgname=labvg --name=home --fstype=xfs --size=1024" >> /tmp/part.ks
    fi
    if [ "$PRESERVE_SRV" = "yes" ]; then
        echo "logvol /srv --vgname=labvg --name=srv --useexisting --noformat" >> /tmp/part.ks
    else
        echo "logvol /srv --vgname=labvg --name=srv --fstype=xfs --size=1024" >> /tmp/part.ks
    fi
    if [ "$PRESERVE_LONGHORN" = "yes" ]; then
        echo "logvol /var/lib/longhorn --vgname=labvg --name=longhorn --useexisting --noformat" >> /tmp/part.ks
    fi
else
    cat > /tmp/part.ks << PARTEOF
ignoredisk --only-use=$DISK
clearpart --all --initlabel --drives=$DISK
part /boot/efi --fstype=efi --size=600 --ondisk=$DISK
part /boot --fstype=ext4 --size=1024 --ondisk=$DISK
part pv.01 --size=1 --grow --ondisk=$DISK
volgroup labvg pv.01
logvol swap --vgname=labvg --name=swap --fstype=swap --size=1024
logvol / --vgname=labvg --name=root --fstype=xfs --size=4096
logvol /var --vgname=labvg --name=var --fstype=xfs --size=3072
logvol /var/log --vgname=labvg --name=varlog --fstype=xfs --size=1024
logvol /home --vgname=labvg --name=home --fstype=xfs --size=1024
logvol /srv --vgname=labvg --name=srv --fstype=xfs --size=1024
logvol /var/lib/longhorn --vgname=labvg --name=longhorn --fstype=xfs --grow --size=1
PARTEOF
fi

echo "=== Generated partition config ==="
cat /tmp/part.ks
%end

%packages
@core
openssh-server
%end

%post
echo "Installed $(date -Iseconds)" > /etc/lab-provisioned
echo "testpass" | passwd --stdin michal
%end
KSEOF
}

# ── Install helper ──
run_install() {
    local label="$1"
    local disk_args="$2"

    log "Running virt-install ($label)..."
    virt-install \
        --name "$VM_NAME" \
        --ram 4096 \
        --vcpus 2 \
        --disk "$disk_args" \
        --os-variant fedora-unknown \
        --network network=default \
        --location "$FEDORA_MIRROR" \
        --initrd-inject "$KS_PATH" \
        --extra-args "inst.ks=file:///test-vm.ks console=ttyS0,115200n8 inst.text" \
        --boot loader="$OVMF_CODE",loader.readonly=yes,loader.type=pflash,nvram.template="$OVMF_VARS" \
        --noautoconsole \
        --wait -1

    log "virt-install exited — install complete"
    virsh destroy "$VM_NAME" 2>/dev/null || true
}

# ── Main test flow ──

generate_kickstart
log "Kickstart generated at $KS_PATH"

PASS=0
FAIL=0

if ! $SKIP_FIRST; then
    # ── Step 1: Fresh install ──
    step "Step 1/4: Fresh install"
    cleanup_all
    run_install "fresh" "path=$DISK_PATH,size=$DISK_SIZE,bus=virtio"

    # Verify fresh install
    log "Verifying fresh install..."
    FILESYSTEMS=$(guestfish --ro -a "$DISK_PATH" -i list-filesystems 2>/dev/null)
    for lv in root var varlog home srv longhorn swap; do
        if echo "$FILESYSTEMS" | grep -q "labvg/$lv"; then
            log "  ✔ labvg/$lv exists"
            ((PASS++))
        else
            err "  ✘ labvg/$lv MISSING"
            ((FAIL++))
        fi
    done
else
    step "Skipping first install (--skip-first-install)"
    [[ -f "$DISK_PATH" ]] || { err "Disk not found at $DISK_PATH"; exit 1; }
fi

# ── Step 2: Write marker files ──
step "Step 2/4: Writing marker files to preserved partitions"
guestfish -a "$DISK_PATH" -i << 'GF'
write /home/michal/PRESERVE_TEST.txt "MARKER: home partition preserved\n"
write /srv/PRESERVE_TEST.txt "MARKER: srv partition preserved\n"
write /var/lib/longhorn/PRESERVE_TEST.txt "MARKER: longhorn partition preserved\n"
write /var/SHOULD_BE_WIPED.txt "This file should NOT survive reprovision\n"
GF
log "Marker files written:"
log "  /home/michal/PRESERVE_TEST.txt"
log "  /srv/PRESERVE_TEST.txt"
log "  /var/lib/longhorn/PRESERVE_TEST.txt"
log "  /var/SHOULD_BE_WIPED.txt (should be wiped)"

# ── Step 3: Reprovision ──
step "Step 3/4: Reprovisioning (reinstall on same disk)"
cleanup_vm
run_install "reprovision" "path=$DISK_PATH,bus=virtio"

# ── Step 4: Verify ──
step "Step 4/4: Verifying preservation"

check_file() {
    local path="$1" expect="$2" label="$3"
    local content
    content=$(guestfish --ro -a "$DISK_PATH" -i cat "$path" 2>/dev/null) || content=""

    if [[ "$expect" == "exists" ]]; then
        if [[ -n "$content" && "$content" == *"MARKER"* ]]; then
            log "  ✔ $label — PRESERVED: $(echo "$content" | head -1)"
            ((PASS++))
        else
            err "  ✘ $label — LOST (file missing or empty)"
            ((FAIL++))
        fi
    elif [[ "$expect" == "gone" ]]; then
        if [[ -z "$content" ]]; then
            log "  ✔ $label — correctly wiped"
            ((PASS++))
        else
            err "  ✘ $label — should have been wiped but still exists"
            ((FAIL++))
        fi
    fi
}

check_file "/home/michal/PRESERVE_TEST.txt" "exists" "/home (preserved)"
check_file "/srv/PRESERVE_TEST.txt" "exists" "/srv (preserved)"
check_file "/var/lib/longhorn/PRESERVE_TEST.txt" "exists" "/var/lib/longhorn (preserved)"
check_file "/var/SHOULD_BE_WIPED.txt" "gone" "/var (wiped)"

# Also verify OS was actually reinstalled
PROV_DATE=$(guestfish --ro -a "$DISK_PATH" -i cat /etc/lab-provisioned 2>/dev/null || echo "")
if [[ -n "$PROV_DATE" ]]; then
    log "  ✔ OS reinstalled: $PROV_DATE"
    ((PASS++))
else
    err "  ✘ /etc/lab-provisioned missing — OS not installed?"
    ((FAIL++))
fi

# ── Summary ──
echo ""
echo -e "${BOLD}════════════════════════════════════════${NC}"
if [[ $FAIL -eq 0 ]]; then
    echo -e "${GREEN}${BOLD}  ALL TESTS PASSED ($PASS/$((PASS+FAIL)))${NC}"
else
    echo -e "${RED}${BOLD}  $FAIL TESTS FAILED ($PASS passed, $FAIL failed)${NC}"
fi
echo -e "${BOLD}════════════════════════════════════════${NC}"
echo ""

# ── Cleanup ──
log "Cleaning up VM (disk preserved at $DISK_PATH)"
cleanup_vm

exit $FAIL