#!/bin/bash
# VM management utilities for archsetup testing
# Author: Craig Jennings <craigmartinjennings@gmail.com>
# License: GNU GPLv3

# Note: logging.sh should already be sourced by the calling script

# VM configuration defaults
VM_CPUS="${VM_CPUS:-4}"
VM_RAM="${VM_RAM:-8192}"  # MB
VM_DISK="${VM_DISK:-50}"  # GB
VM_NETWORK="${VM_NETWORK:-default}"
LIBVIRT_URI="qemu:///system"  # Use system session, not user session

# Check if libvirt is running
check_libvirt() {
    if ! systemctl is-active --quiet libvirtd; then
        error "libvirtd service is not running"
        info "Start it with: sudo systemctl start libvirtd"
        return 1
    fi
    return 0
}

# Check if user is in libvirt group
check_libvirt_group() {
    if ! groups | grep -q libvirt; then
        warn "Current user is not in libvirt group"
        info "Add yourself with: sudo usermod -a -G libvirt $USER"
        info "Then log out and back in for changes to take effect"
        return 1
    fi
    return 0
}

# Check if KVM is available
check_kvm() {
    if [ ! -e /dev/kvm ]; then
        error "KVM is not available"
        info "Check if virtualization is enabled in BIOS"
        info "Load kvm module: sudo modprobe kvm-intel (or kvm-amd)"
        return 1
    fi
    return 0
}

# Wait for VM to boot (check for SSH or serial console)
wait_for_vm() {
    local vm_name="$1"
    local timeout="${2:-300}"  # 5 minutes default
    local elapsed=0

    progress "Waiting for VM $vm_name to boot..."

    while [ $elapsed -lt $timeout ]; do
        if virsh --connect "$LIBVIRT_URI" domstate "$vm_name" 2>/dev/null | grep -q "running"; then
            sleep 5
            complete "VM $vm_name is running"
            return 0
        fi
        sleep 2
        elapsed=$((elapsed + 2))
    done

    error "Timeout waiting for VM $vm_name to boot"
    return 1
}

# Check if VM exists
vm_exists() {
    local vm_name="$1"
    virsh --connect "$LIBVIRT_URI" dominfo "$vm_name" &>/dev/null
    return $?
}

# Check if VM is running
vm_is_running() {
    local vm_name="$1"
    [ "$(virsh --connect "$LIBVIRT_URI" domstate "$vm_name" 2>/dev/null)" = "running" ]
    return $?
}

# Start VM
start_vm() {
    local vm_name="$1"

    if vm_is_running "$vm_name"; then
        warn "VM $vm_name is already running"
        return 0
    fi

    step "Starting VM: $vm_name"
    if virsh --connect "$LIBVIRT_URI" start "$vm_name" >> "$LOGFILE" 2>&1; then
        success "VM $vm_name started"
        return 0
    else
        error "Failed to start VM $vm_name"
        return 1
    fi
}

# Stop VM gracefully
stop_vm() {
    local vm_name="$1"
    local timeout="${2:-60}"

    if ! vm_is_running "$vm_name"; then
        info "VM $vm_name is not running"
        return 0
    fi

    step "Shutting down VM: $vm_name"
    if virsh --connect "$LIBVIRT_URI" shutdown "$vm_name" >> "$LOGFILE" 2>&1; then
        # Wait for graceful shutdown
        local elapsed=0
        while [ $elapsed -lt $timeout ]; do
            if ! vm_is_running "$vm_name"; then
                success "VM $vm_name stopped gracefully"
                return 0
            fi
            sleep 2
            elapsed=$((elapsed + 2))
        done

        warn "VM $vm_name did not stop gracefully, forcing..."
        virsh --connect "$LIBVIRT_URI" destroy "$vm_name" >> "$LOGFILE" 2>&1
    fi

    success "VM $vm_name stopped"
    return 0
}

# Destroy VM (force stop)
destroy_vm() {
    local vm_name="$1"

    if ! vm_exists "$vm_name"; then
        info "VM $vm_name does not exist"
        return 0
    fi

    step "Destroying VM: $vm_name"
    if vm_is_running "$vm_name"; then
        virsh --connect "$LIBVIRT_URI" destroy "$vm_name" >> "$LOGFILE" 2>&1
    fi

    virsh --connect "$LIBVIRT_URI" undefine "$vm_name" --nvram >> "$LOGFILE" 2>&1
    success "VM $vm_name destroyed"
    return 0
}

# Create snapshot
create_snapshot() {
    local vm_name="$1"
    local snapshot_name="$2"

    step "Creating snapshot: $snapshot_name"
    if virsh --connect "$LIBVIRT_URI" snapshot-create-as "$vm_name" "$snapshot_name" >> "$LOGFILE" 2>&1; then
        success "Snapshot $snapshot_name created"
        return 0
    else
        error "Failed to create snapshot $snapshot_name"
        return 1
    fi
}

# Restore snapshot
restore_snapshot() {
    local vm_name="$1"
    local snapshot_name="$2"

    step "Restoring snapshot: $snapshot_name"
    if virsh --connect "$LIBVIRT_URI" snapshot-revert "$vm_name" "$snapshot_name" >> "$LOGFILE" 2>&1; then
        success "Snapshot $snapshot_name restored"
        return 0
    else
        error "Failed to restore snapshot $snapshot_name"
        return 1
    fi
}

# Delete snapshot
delete_snapshot() {
    local vm_name="$1"
    local snapshot_name="$2"

    step "Deleting snapshot: $snapshot_name"
    if virsh --connect "$LIBVIRT_URI" snapshot-delete "$vm_name" "$snapshot_name" >> "$LOGFILE" 2>&1; then
        success "Snapshot $snapshot_name deleted"
        return 0
    else
        error "Failed to delete snapshot $snapshot_name"
        return 1
    fi
}

# Clone disk image (copy-on-write)
clone_disk() {
    local base_image="$1"
    local new_image="$2"

    if [ ! -f "$base_image" ]; then
        error "Base image not found: $base_image"
        return 1
    fi

    step "Cloning disk image (full copy)"
    if qemu-img convert -f qcow2 -O qcow2 "$base_image" "$new_image" >> "$LOGFILE" 2>&1; then
        success "Disk cloned: $new_image"
    else
        error "Failed to clone disk"
        return 1
    fi

    # Truncate machine-id so systemd generates a new one on boot (avoids DHCP conflicts)
    step "Clearing machine-id for unique network identity"
    if guestfish -a "$new_image" -i truncate /etc/machine-id >> "$LOGFILE" 2>&1; then
        success "Machine-ID cleared (will regenerate on boot)"
        return 0
    else
        warn "Failed to clear machine-ID (guestfish failed)"
        info "Network may conflict with base VM if both run simultaneously"
        return 0  # Don't fail the whole operation
    fi
}

# Get VM IP address (requires guest agent or DHCP lease)
get_vm_ip() {
    local vm_name="$1"

    # Try guest agent first
    local ip
    ip=$(virsh --connect "$LIBVIRT_URI" domifaddr "$vm_name" 2>/dev/null | grep -oP '(\d+\.){3}\d+' | head -1)

    if [ -n "$ip" ]; then
        echo "$ip"
        return 0
    fi

    # Fall back to DHCP leases
    local mac
    mac=$(virsh --connect "$LIBVIRT_URI" domiflist "$vm_name" | grep -oP '([0-9a-f]{2}:){5}[0-9a-f]{2}' | head -1)

    if [ -n "$mac" ]; then
        ip=$(grep "$mac" /var/lib/libvirt/dnsmasq/default.leases 2>/dev/null | awk '{print $3}')
        if [ -n "$ip" ]; then
            echo "$ip"
            return 0
        fi
    fi

    return 1
}

# Execute command in VM via SSH
vm_exec() {
    local vm_name="$1"
    shift
    local cmd="$*"

    local ip
    ip=$(get_vm_ip "$vm_name")

    if [ -z "$ip" ]; then
        error "Could not get IP address for VM $vm_name"
        return 1
    fi

    ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
        "root@$ip" "$cmd" 2>> "$LOGFILE"
}

# Copy file to VM
copy_to_vm() {
    local vm_name="$1"
    local local_file="$2"
    local remote_path="$3"

    local ip
    ip=$(get_vm_ip "$vm_name")

    if [ -z "$ip" ]; then
        error "Could not get IP address for VM $vm_name"
        return 1
    fi

    step "Copying $local_file to VM"
    if scp -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
        "$local_file" "root@$ip:$remote_path" >> "$LOGFILE" 2>&1; then
        success "File copied to VM"
        return 0
    else
        error "Failed to copy file to VM"
        return 1
    fi
}

# Copy file from VM
copy_from_vm() {
    local vm_name="$1"
    local remote_file="$2"
    local local_path="$3"

    local ip
    ip=$(get_vm_ip "$vm_name")

    if [ -z "$ip" ]; then
        error "Could not get IP address for VM $vm_name"
        return 1
    fi

    step "Copying $remote_file from VM"
    if scp -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
        "root@$ip:$remote_file" "$local_path" >> "$LOGFILE" 2>&1; then
        success "File copied from VM"
        return 0
    else
        error "Failed to copy file from VM"
        return 1
    fi
}