#!/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. Boots the VM via QEMU
# 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"
source "$SCRIPT_DIR/lib/validation.sh"

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

while [[ $# -gt 0 ]]; do
    case $1 in
        --keep)
            KEEP_VM=true
            shift
            ;;
        --script)
            ARCHSETUP_SCRIPT="$2"
            shift 2
            ;;
        --snapshot)
            SNAPSHOT_NAME="$2"
            shift 2
            ;;
        *)
            echo "Usage: $0 [--keep] [--script /path/to/archsetup] [--snapshot name]"
            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)"
            exit 1
            ;;
    esac
done

# Configuration
TIMESTAMP=$(date +'%Y%m%d-%H%M%S')
VM_IMAGES_DIR="$PROJECT_ROOT/vm-images"
TEST_RESULTS_DIR="$PROJECT_ROOT/test-results/$TIMESTAMP"
ROOT_PASSWORD="archsetup"
ARCHZFS_INBOX="$HOME/code/archzfs/inbox"
ARCHSETUP_VM_CONF="$SCRIPT_DIR/archsetup-vm.conf"

# Set VM_IP for validation.sh backward compatibility
VM_IP="localhost"

# Initialize logging and VM paths
mkdir -p "$TEST_RESULTS_DIR"
LOGFILE="$TEST_RESULTS_DIR/test.log"
init_logging "$LOGFILE"
init_vm_paths "$VM_IMAGES_DIR"

section "ArchSetup Test Run: $TIMESTAMP"

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

# Check disk exists
if [ ! -f "$DISK_PATH" ]; then
    fatal "Base disk not found: $DISK_PATH"
    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 ! snapshot_exists "$DISK_PATH" "$SNAPSHOT_NAME"; then
    fatal "Snapshot '$SNAPSHOT_NAME' not found on $DISK_PATH"
    info "Available snapshots:"
    list_snapshots "$DISK_PATH"
    info ""
    info "Create base VM with: ./scripts/testing/create-base-vm.sh"
fi
success "Snapshot $SNAPSHOT_NAME exists"

# Stop VM if running and restore snapshot
stop_qemu 2>/dev/null || true

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

# Start VM and wait for SSH
start_timer "boot"
step "Starting VM and waiting for SSH..."
start_qemu "$DISK_PATH" "disk" "" "none" || fatal "Failed to start VM"
wait_for_ssh "$ROOT_PASSWORD" 120 || fatal "VM SSH not available"
stop_timer "boot"

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

# Capture pre-install state for comparison
capture_pre_install_state "$TEST_RESULTS_DIR"

# 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
vm_exec "$ROOT_PASSWORD" "rm -rf /tmp/archsetup-test && mkdir -p /tmp/archsetup-test" >> "$LOGFILE" 2>&1

copy_to_vm "$BUNDLE_FILE" "/tmp/archsetup.bundle" "$ROOT_PASSWORD"

# Clone from bundle on VM (simulates git clone)
vm_exec "$ROOT_PASSWORD" \
    "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)"

# Transfer archsetup VM config file
step "Copying archsetup VM config"
copy_to_vm "$ARCHSETUP_VM_CONF" "/tmp/archsetup-test/archsetup-vm.conf" "$ROOT_PASSWORD"

# Pre-archsetup VM setup
section "Pre-ArchSetup VM Setup"

step "Importing archzfs PGP key"
vm_exec "$ROOT_PASSWORD" \
    "curl -sL https://archzfs.com/archzfs.gpg | pacman-key --add - && pacman-key --lsign-key DDF7DB817396A49B2A2723F7403BD972F75D9D76" \
    >> "$LOGFILE" 2>&1 && success "archzfs PGP key imported" || warn "Failed to import archzfs key (may already be present)"

step "Syncing package databases"
vm_exec "$ROOT_PASSWORD" "pacman -Sy --noconfirm" >> "$LOGFILE" 2>&1 && \
    success "Package databases synced" || warn "Package database sync had issues"

# Execute archsetup
section "Executing ArchSetup"

start_timer "archsetup"
step "Starting archsetup script in detached session on VM..."
info "Log file: $LOGFILE"

# Start archsetup fully detached.
# Use ssh -T -n to prevent PTY allocation and stdin forwarding, which allows
# the SSH session to close immediately after the remote shell exits.
REMOTE_LOG="/tmp/archsetup-test/archsetup-output.log"

sshpass -p "$ROOT_PASSWORD" ssh -T -n $SSH_OPTS \
    -p "$SSH_PORT" root@localhost \
    "setsid bash -c 'cd /tmp/archsetup-test && bash archsetup --config-file /tmp/archsetup-test/archsetup-vm.conf > $REMOTE_LOG 2>&1' < /dev/null > /dev/null 2>&1 &" \
    >> "$LOGFILE" 2>&1

# Verify the process started
sleep 3
if vm_exec "$ROOT_PASSWORD" "pgrep -f 'bash archsetup'" >> "$LOGFILE" 2>/dev/null; then
    success "ArchSetup started in background on VM"
else
    fatal "ArchSetup process not found after launch"
fi

# 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
    if vm_exec "$ROOT_PASSWORD" "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=$(vm_exec "$ROOT_PASSWORD" \
        "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..."
copy_from_vm "$REMOTE_LOG" "$TEST_RESULTS_DIR/archsetup-output.log" "$ROOT_PASSWORD" || \
    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"
copy_from_vm "/var/log/archsetup-*.log" "$TEST_RESULTS_DIR/" "$ROOT_PASSWORD" || \
    warn "Could not copy archsetup log"

step "Copying package lists from VM"
copy_from_vm "/var/log/archsetup-*-package-list.txt" "$TEST_RESULTS_DIR/" "$ROOT_PASSWORD" || \
    warn "Could not copy package lists"

step "Copying installed packages list"
copy_from_vm "/var/log/archsetup-installed-packages.txt" "$TEST_RESULTS_DIR/" "$ROOT_PASSWORD" || \
    warn "Could not copy installed packages list"

# Capture post-install state
capture_post_install_state "$TEST_RESULTS_DIR"

# Run comprehensive validation
# This uses the validation.sh library for all checks
run_all_validations
validate_all_services

# Analyze log differences (pre vs post install)
analyze_log_diff "$TEST_RESULTS_DIR"

# Generate issue attribution report
# If base install issues found and archzfs inbox exists, create issue file
generate_issue_report "$TEST_RESULTS_DIR" "$ARCHZFS_INBOX"

# Set validation result based on failure count
if [ $VALIDATION_FAILED -eq 0 ]; then
    TEST_PASSED=true
else
    TEST_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: QEMU snapshot-based

VM Configuration:
  Disk:     $DISK_PATH
  Snapshot: $SNAPSHOT_NAME
  SSH:      localhost:$SSH_PORT

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

Validation Summary:
  Passed:   $VALIDATION_PASSED
  Failed:   $VALIDATION_FAILED
  Warnings: $VALIDATION_WARNINGS

Issue Attribution:
  ArchSetup Issues: ${#ARCHSETUP_ISSUES[@]}
  Base Install Issues: ${#BASE_INSTALL_ISSUES[@]}
  Unknown Issues: ${#UNKNOWN_ISSUES[@]}

Artifacts:
  Log file: $LOGFILE
  Report:   $REPORT_FILE
  Results:  $TEST_RESULTS_DIR/
  Issue Report: $TEST_RESULTS_DIR/issue-report.txt
  Pre-install logs: $TEST_RESULTS_DIR/pre-install/
  Post-install logs: $TEST_RESULTS_DIR/post-install/
  Analysis: $TEST_RESULTS_DIR/analysis/

EOFREPORT

info "Test report saved: $REPORT_FILE"

# Display report to terminal
echo ""
cat "$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 "  SSH: sshpass -p '$ROOT_PASSWORD' ssh -p $SSH_PORT root@localhost"
    info ""
    info "To stop VM: kill \$(cat $PID_FILE)"
    info "To revert:  qemu-img snapshot -a $SNAPSHOT_NAME $DISK_PATH"
else
    step "Shutting down VM and reverting to clean snapshot"
    stop_qemu
    if restore_snapshot "$DISK_PATH" "$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" = "0" ] && $TEST_PASSED; then
    success "TEST PASSED"
    exit 0
else
    error "TEST FAILED"
    info "Check logs in: $TEST_RESULTS_DIR"
    exit 1
fi