#!/bin/bash
# Clean up old test VMs and artifacts
# Author: Craig Jennings <craigmartinjennings@gmail.com>
# License: GNU GPLv3

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"

# Parse arguments
KEEP_LAST=5
FORCE=false

while [[ $# -gt 0 ]]; do
    case $1 in
        --keep)
            KEEP_LAST="$2"
            shift 2
            ;;
        --force)
            FORCE=true
            shift
            ;;
        *)
            echo "Usage: $0 [--keep N] [--force]"
            echo "  --keep N  Keep last N test results (default: 5)"
            echo "  --force   Skip confirmation prompts"
            exit 1
            ;;
    esac
done

# Initialize logging and VM paths
LOGFILE="/tmp/cleanup-tests-$(date +'%Y%m%d-%H%M%S').log"
init_logging "$LOGFILE"
init_vm_paths "$PROJECT_ROOT/vm-images"

section "Cleaning Up Test Artifacts"

# Find and stop running QEMU processes
step "Checking for running QEMU processes"

if vm_is_running; then
    info "Found running QEMU test VM (PID: $(cat "$PID_FILE"))"
    if $FORCE; then
        stop_qemu
    else
        read -p "Stop running VM? [y/N] " -n 1 -r
        echo ""
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            stop_qemu
        else
            info "Skipping VM shutdown"
        fi
    fi
else
    info "No running VM found"
fi

# Check for orphaned QEMU processes
QEMU_PIDS=$(pgrep -f "qemu-system.*archsetup-test" 2>/dev/null || true)
if [ -n "$QEMU_PIDS" ]; then
    info "Found orphaned QEMU processes: $QEMU_PIDS"
    if $FORCE; then
        echo "$QEMU_PIDS" | xargs kill -9 2>/dev/null || true
        success "Orphaned processes killed"
    else
        read -p "Kill orphaned QEMU processes? [y/N] " -n 1 -r
        echo ""
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            echo "$QEMU_PIDS" | xargs kill -9 2>/dev/null || true
            success "Orphaned processes killed"
        fi
    fi
fi

# Clean up QEMU runtime files
rm -f "$PID_FILE" "$MONITOR_SOCK"

# Clean up debug overlay and test disk images
section "Cleaning Up Disk Images"

step "Finding temporary disk images"
if [ -d "$PROJECT_ROOT/vm-images" ]; then
    TEMP_DISKS=$(find "$PROJECT_ROOT/vm-images" -name "debug-overlay-*.qcow2" -o -name "archsetup-test-*.qcow2" 2>/dev/null || true)

    if [ -z "$TEMP_DISKS" ]; then
        info "No temporary disk images found"
    else
        DISK_COUNT=$(echo "$TEMP_DISKS" | wc -l)
        DISK_SIZE=$(du -ch $TEMP_DISKS 2>/dev/null | tail -1 | awk '{print $1}')
        info "Found $DISK_COUNT temporary disk image(s) totaling $DISK_SIZE"

        if $FORCE; then
            echo "$TEMP_DISKS" | while read disk; do
                rm -f "$disk"
            done
            success "Temporary disk images deleted"
        else
            echo ""
            echo "$TEMP_DISKS"
            echo ""
            read -p "Delete these disk images? [y/N] " -n 1 -r
            echo ""
            if [[ $REPLY =~ ^[Yy]$ ]]; then
                echo "$TEMP_DISKS" | while read disk; do
                    rm -f "$disk"
                done
                success "Temporary disk images deleted"
            else
                info "Skipping disk cleanup"
            fi
        fi
    fi
fi

# Clean up old test results
section "Cleaning Up Test Results"

if [ ! -d "$PROJECT_ROOT/test-results" ]; then
    info "No test results directory"
else
    step "Finding test result directories"
    TEST_RESULTS=$(find "$PROJECT_ROOT/test-results" -maxdepth 1 -type d -name "20*" 2>/dev/null | sort -r || true)

    if [ -z "$TEST_RESULTS" ]; then
        info "No test results found"
    else
        RESULT_COUNT=$(echo "$TEST_RESULTS" | wc -l)
        info "Found $RESULT_COUNT test result directory(ies)"

        if [ $RESULT_COUNT -le $KEEP_LAST ]; then
            info "Keeping all results (count <= $KEEP_LAST)"
        else
            TO_DELETE=$(echo "$TEST_RESULTS" | tail -n +$((KEEP_LAST + 1)))
            DELETE_COUNT=$(echo "$TO_DELETE" | wc -l)
            info "Keeping last $KEEP_LAST, deleting $DELETE_COUNT old result(s)"

            if $FORCE; then
                echo "$TO_DELETE" | while read dir; do
                    rm -rf "$dir"
                done
                success "Old test results deleted"
            else
                echo ""
                echo "Will delete:"
                echo "$TO_DELETE"
                echo ""
                read -p "Delete these test results? [y/N] " -n 1 -r
                echo ""
                if [[ $REPLY =~ ^[Yy]$ ]]; then
                    echo "$TO_DELETE" | while read dir; do
                        rm -rf "$dir"
                    done
                    success "Old test results deleted"
                else
                    info "Skipping results cleanup"
                fi
            fi
        fi
    fi
fi

section "Cleanup Complete"

info "Log file: $LOGFILE"