#!/bin/bash
# Run archsetup test in a VM using snapshots
# Author: Craig Jennings <craigmartinjennings@gmail.com>
# License: GNU GPLv3
#
# This script:
# 1. Reverts base VM to clean snapshot
# 2. Starts the base VM
# 3. Transfers archsetup and dotfiles
# 4. Executes archsetup in the VM
# 5. Captures logs and validates results
# 6. Generates test report
# 7. Reverts to clean snapshot for next run

set -e

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

# Source utilities
source "$SCRIPT_DIR/lib/logging.sh"
source "$SCRIPT_DIR/lib/vm-utils.sh"
source "$SCRIPT_DIR/lib/network-diagnostics.sh"

# Parse arguments
KEEP_VM=false
ARCHSETUP_SCRIPT="$PROJECT_ROOT/archsetup"
SNAPSHOT_NAME="clean-install"
SKIP_SLOW_PACKAGES=false

while [[ $# -gt 0 ]]; do
    case $1 in
        --keep)
            KEEP_VM=true
            shift
            ;;
        --script)
            ARCHSETUP_SCRIPT="$2"
            shift 2
            ;;
        --snapshot)
            SNAPSHOT_NAME="$2"
            shift 2
            ;;
        --skip-slow-packages)
            SKIP_SLOW_PACKAGES=true
            shift
            ;;
        *)
            echo "Usage: $0 [--keep] [--script /path/to/archsetup] [--snapshot name] [--skip-slow-packages]"
            echo "  --keep                Keep VM in post-test state (for debugging)"
            echo "  --script              Specify custom archsetup script to test"
            echo "  --snapshot            Snapshot name to revert to (default: clean-install)"
            echo "  --skip-slow-packages  Skip slow packages (texlive-meta, topgrade) for faster testing"
            exit 1
            ;;
    esac
done

# Configuration
TIMESTAMP=$(date +'%Y%m%d-%H%M%S')
VM_NAME="archsetup-base"
TEST_RESULTS_DIR="$PROJECT_ROOT/test-results/$TIMESTAMP"
ROOT_PASSWORD="archsetup"

# Initialize logging
mkdir -p "$TEST_RESULTS_DIR"
LOGFILE="$TEST_RESULTS_DIR/test.log"
init_logging "$LOGFILE"

section "ArchSetup Test Run: $TIMESTAMP"

# Verify archsetup script exists
if [ ! -f "$ARCHSETUP_SCRIPT" ]; then
    fatal "ArchSetup script not found: $ARCHSETUP_SCRIPT"
fi

# Check if VM exists
if ! vm_exists "$VM_NAME"; then
    fatal "Base VM not found: $VM_NAME"
    info "Create it first: ./scripts/testing/create-base-vm.sh"
fi

# Check if snapshot exists
section "Preparing Test Environment"

step "Checking for snapshot: $SNAPSHOT_NAME"
if ! virsh --connect "$LIBVIRT_URI" snapshot-list "$VM_NAME" --name 2>/dev/null | grep -q "^$SNAPSHOT_NAME$"; then
    fatal "Snapshot '$SNAPSHOT_NAME' not found on VM $VM_NAME"
    info "Available snapshots:"
    virsh --connect "$LIBVIRT_URI" snapshot-list "$VM_NAME" 2>/dev/null || info "  (none)"
    info ""
    info "Create snapshot with:"
    info "  virsh snapshot-create-as $VM_NAME $SNAPSHOT_NAME --description 'Clean Arch install'"
fi
success "Snapshot $SNAPSHOT_NAME exists"

# Shut down VM if running
if vm_is_running "$VM_NAME"; then
    warn "VM $VM_NAME is currently running - shutting down for snapshot revert"
    stop_vm "$VM_NAME"
fi

# Revert to clean snapshot
step "Reverting to snapshot: $SNAPSHOT_NAME"
if restore_snapshot "$VM_NAME" "$SNAPSHOT_NAME"; then
    success "Reverted to clean state"
else
    fatal "Failed to revert snapshot"
fi

# Start VM
start_timer "boot"
step "Starting VM and waiting for SSH..."
if ! start_vm "$VM_NAME"; then
    fatal "Failed to start VM"
fi

sleep 10  # Give VM time to boot

# Get VM IP address
VM_IP=""
for i in {1..30}; do
    VM_IP=$(get_vm_ip "$VM_NAME" 2>/dev/null || true)
    if [ -n "$VM_IP" ]; then
        break
    fi
    sleep 2
done

if [ -z "$VM_IP" ]; then
    error "Could not get VM IP address"
    info "VM may not have booted correctly"
    fatal "VM boot failed"
fi

success "VM is running at $VM_IP"
stop_timer "boot"

# Wait for SSH
step "Waiting for SSH to become available..."
for i in {1..60}; do
    if sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
           -o ConnectTimeout=2 "root@$VM_IP" "echo connected" &>/dev/null; then
        break
    fi
    sleep 2
done

success "SSH is available"

# Run network diagnostics
if ! run_network_diagnostics "$VM_IP"; then
    fatal "Network diagnostics failed - aborting test"
fi

# Transfer files to VM (simulating git clone)
section "Simulating Git Clone"

step "Creating shallow git clone on VM"
info "This simulates: git clone --depth 1 <repo> /home/cjennings/code/archsetup"

# Create a temporary git bundle from current repo
BUNDLE_FILE=$(mktemp)
git bundle create "$BUNDLE_FILE" HEAD >> "$LOGFILE" 2>&1

# Transfer bundle and extract on VM
sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    "root@$VM_IP" "rm -rf /tmp/archsetup-test && mkdir -p /tmp/archsetup-test" >> "$LOGFILE" 2>&1

sshpass -p "$ROOT_PASSWORD" scp -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    "$BUNDLE_FILE" "root@$VM_IP:/tmp/archsetup.bundle" >> "$LOGFILE" 2>&1

# Clone from bundle on VM (simulates git clone)
sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    "root@$VM_IP" "cd /tmp && git clone --depth 1 /tmp/archsetup.bundle archsetup-test && rm /tmp/archsetup.bundle" >> "$LOGFILE" 2>&1

rm -f "$BUNDLE_FILE"
success "Repository cloned to VM (simulating git clone --depth 1)"

# Execute archsetup
section "Executing ArchSetup"

start_timer "archsetup"
step "Starting archsetup script in detached session on VM..."
info "This will take 30-60 minutes depending on network speed"
info "Log file: $LOGFILE"

# Start archsetup in a detached session on the VM (resilient to SSH disconnections)
REMOTE_LOG="/tmp/archsetup-test/archsetup-output.log"
ARCHSETUP_ARGS=""
if $SKIP_SLOW_PACKAGES; then
    ARCHSETUP_ARGS="--skip-slow-packages"
    info "Running archsetup with --skip-slow-packages flag"
fi
sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    "root@$VM_IP" "cd /tmp/archsetup-test && nohup bash archsetup $ARCHSETUP_ARGS > $REMOTE_LOG 2>&1 & echo \$!" \
    >> "$LOGFILE" 2>&1

if [ $? -ne 0 ]; then
    fatal "Failed to start archsetup on VM"
fi

success "ArchSetup started in background on VM"

# Poll for completion
step "Monitoring archsetup progress (polling every 30 seconds)..."
POLL_COUNT=0
MAX_POLLS=180  # 90 minutes max (180 * 30 seconds)

while [ $POLL_COUNT -lt $MAX_POLLS ]; do
    # Check if archsetup process is still running
    # Use ps to avoid pgrep matching its own SSH command
    if sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
           "root@$VM_IP" "ps aux | grep '[b]ash archsetup' > /dev/null" 2>/dev/null; then
        # Still running, wait and continue
        sleep 30
        POLL_COUNT=$((POLL_COUNT + 1))

        # Show progress every 5 minutes
        if [ $((POLL_COUNT % 10)) -eq 0 ]; then
            ELAPSED_MINS=$((POLL_COUNT / 2))
            info "Still running... ($ELAPSED_MINS minutes elapsed)"
        fi
    else
        # Process finished
        break
    fi
done

if [ $POLL_COUNT -ge $MAX_POLLS ]; then
    error "ArchSetup timed out after 90 minutes"
    ARCHSETUP_EXIT_CODE=124
else
    # Get exit code from the remote log
    step "Retrieving archsetup exit status..."
    ARCHSETUP_EXIT_CODE=$(sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
        "root@$VM_IP" "grep -q 'ARCHSETUP_EXECUTION_COMPLETE' /var/log/archsetup-*.log 2>/dev/null && echo 0 || echo 1" 2>/dev/null)

    if [ "$ARCHSETUP_EXIT_CODE" = "0" ]; then
        success "ArchSetup completed successfully"
    else
        error "ArchSetup may have encountered errors (check logs)"
    fi
fi

# Copy the remote output log
step "Retrieving archsetup output from VM..."
sshpass -p "$ROOT_PASSWORD" scp -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    "root@$VM_IP:$REMOTE_LOG" "$TEST_RESULTS_DIR/archsetup-output.log" 2>> "$LOGFILE" || \
    warn "Could not copy remote output log"

# Append remote output to main test log
if [ -f "$TEST_RESULTS_DIR/archsetup-output.log" ]; then
    cat "$TEST_RESULTS_DIR/archsetup-output.log" >> "$LOGFILE"
fi

stop_timer "archsetup"

# Capture logs and artifacts from VM
section "Capturing Test Artifacts"

step "Copying archsetup log from VM"
sshpass -p "$ROOT_PASSWORD" scp -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    "root@$VM_IP:/var/log/archsetup-*.log" "$TEST_RESULTS_DIR/" 2>> "$LOGFILE" || \
    warn "Could not copy archsetup log"

step "Copying package lists from VM"
sshpass -p "$ROOT_PASSWORD" scp -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
    "root@$VM_IP:/root/.local/src/archsetup-*.txt" "$TEST_RESULTS_DIR/" 2>> "$LOGFILE" || \
    warn "Could not copy package lists"

# Run validation
section "Validating Installation"

VALIDATION_PASSED=true

# Check if user was created
step "Checking if user 'cjennings' was created"
if sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
       "root@$VM_IP" "id cjennings" &>> "$LOGFILE"; then
    success "User cjennings exists"
else
    error "User cjennings not found"
    VALIDATION_PASSED=false
fi

# Check if dotfiles were stowed
step "Checking if dotfiles are stowed"
if sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
       "root@$VM_IP" "test -L /home/cjennings/.profile" &>> "$LOGFILE"; then
    success "Dotfiles appear to be stowed"
else
    warn "Dotfiles may not be properly stowed"
    VALIDATION_PASSED=false
fi

# Check if yay is installed
step "Checking if yay (AUR helper) is installed"
if sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
       "root@$VM_IP" "/usr/bin/which yay" &>> "$LOGFILE"; then
    success "yay is installed"
else
    error "yay not found"
    VALIDATION_PASSED=false
fi

# Check if DWM was built
step "Checking if DWM is installed"
if sshpass -p "$ROOT_PASSWORD" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
       "root@$VM_IP" "test -f /usr/local/bin/dwm" &>> "$LOGFILE"; then
    success "DWM is installed"
else
    error "DWM not found"
    VALIDATION_PASSED=false
fi

# Generate test report
section "Generating Test Report"

REPORT_FILE="$TEST_RESULTS_DIR/test-report.txt"
cat > "$REPORT_FILE" << EOFREPORT
========================================
ArchSetup Test Report
========================================

Test ID: $TIMESTAMP
Date: $(date +'%Y-%m-%d %H:%M:%S')
Test Method: Snapshot-based

VM Configuration:
  Name: $VM_NAME
  IP:   $VM_IP
  Snapshot: $SNAPSHOT_NAME

Results:
  ArchSetup Exit Code: $ARCHSETUP_EXIT_CODE
  Validation: $(if $VALIDATION_PASSED; then echo "PASSED"; else echo "FAILED"; fi)

Artifacts:
  Log file: $LOGFILE
  Report:   $REPORT_FILE
  Results:  $TEST_RESULTS_DIR/

EOFREPORT

info "Test report saved: $REPORT_FILE"

# Cleanup or keep VM
section "Cleanup"

if $KEEP_VM; then
    info "VM is still running in post-test state (--keep flag was used)"
    info "Connect with:"
    info "  Console: virsh console $VM_NAME"
    info "  SSH:     ssh root@$VM_IP"
    info ""
    info "To revert to clean state when done:"
    info "  virsh shutdown $VM_NAME"
    info "  virsh snapshot-revert $VM_NAME $SNAPSHOT_NAME"
else
    step "Shutting down VM and reverting to clean snapshot"
    stop_vm "$VM_NAME"
    if restore_snapshot "$VM_NAME" "$SNAPSHOT_NAME"; then
        success "VM reverted to clean state"
    else
        warn "Failed to revert snapshot - VM may be in modified state"
    fi
fi

# Final summary
section "Test Complete"

if [ $ARCHSETUP_EXIT_CODE -eq 0 ] && $VALIDATION_PASSED; then
    success "TEST PASSED"
    exit 0
else
    error "TEST FAILED"
    info "Check logs in: $TEST_RESULTS_DIR"
    exit 1
fi