#!/bin/bash
set -e

# Configuration
VM_DIR="${VM_DIR:-$(pwd)/control-plane-vm}"
CONTROLPLANE_NAME="talos-controlplane"
MEMORY="2048"  # 2GB per VM
CPUS="2"
DISK_SIZE="10G"  # 10GB disk
TALOS_ISO_URL="https://github.com/siderolabs/talos/releases/download/v1.11.0/metal-arm64.iso"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Create VM directory
create_vm_dir() {
    local vm_name=$1
    local vm_path="$VM_DIR/$vm_name"
    
    mkdir -p "$vm_path"
    echo "$vm_path"
}

# Download Talos ISO if not exists
download_talos_iso() {
    local iso_path="$VM_DIR/talos.iso"
    
    if [[ ! -f "$iso_path" ]]; then
        log_info "Downloading Talos ISO..."
        curl -L -o "$iso_path" "$TALOS_ISO_URL"
    else
        log_info "Talos ISO already exists"
    fi
    echo "$iso_path"
}

# Create disk image
create_disk_image() {
    local vm_path=$1
    local vm_name=$2
    local disk_path="$vm_path/disk.qcow2"
    
    if [[ ! -f "$disk_path" ]]; then
        log_info "Creating disk image for $vm_name..."
        qemu-img create -f qcow2 "$disk_path" "$DISK_SIZE"
    else
        log_warn "Disk image already exists for $vm_name"
    fi
    echo "$disk_path"
}

# Generate unique MAC address
generate_mac() {
    local vm_name=$1
    local hash=$(echo "$vm_name" | md5sum | cut -c1-6)
    printf "52:54:00:%s:%s:%s" "${hash:0:2}" "${hash:2:2}" "${hash:4:2}"
}

# Create VM startup script
create_vm_script() {
    local vm_path=$1
    local vm_name=$2
    local mac_address=$3
    local iso_path=$4
    local disk_path=$5
    
    local script_path="$vm_path/start.sh"
    
    cat > "$script_path" << 'EOF'
#!/bin/bash
VM_NAME="$1"
MAC_ADDRESS="$2"
ISO_PATH="$3"
DISK_PATH="$4"
CPUS="$5"
MEMORY="$6"

qemu-system-aarch64 \
    -name "$VM_NAME" \
    -machine virt,highmem=off \
    -accel hvf \
    -cpu host \
    -smp "$CPUS" \
    -m "${MEMORY}M" \
    -drive file="$DISK_PATH",if=virtio,format=qcow2 \
    -cdrom "$ISO_PATH" \
    -netdev user,id=net0 \
    -device virtio-net-pci,netdev=net0,mac="$MAC_ADDRESS" \
    -nographic \
    -serial mon:stdio
EOF
    
    # Make the script executable and create a wrapper that passes the parameters
    chmod +x "$script_path"
    
    # Create a wrapper script that calls the main script with proper arguments
    local wrapper_path="$vm_path/run.sh"
    cat > "$wrapper_path" << EOF
#!/bin/bash
"$script_path" "$vm_name" "$mac_address" "$iso_path" "$disk_path" "$CPUS" "$MEMORY"
EOF
    
    chmod +x "$wrapper_path"
    echo "$wrapper_path"
}

# Create systemd service file (optional)
create_service_file() {
    local vm_path=$1
    local vm_name=$2
    local script_path=$3
    
    local service_path="$vm_path/$vm_name.service"
    
    cat > "$service_path" << EOF
[Unit]
Description=Talos VM - $vm_name
After=network.target

[Service]
Type=simple
ExecStart=$script_path
WorkingDirectory=$vm_path
Restart=always
User=$USER

[Install]
WantedBy=multi-user.target
EOF
    
    echo "$service_path"
}

# Main execution
main() {
    log_info "Creating Talos VMs with QEMU..."
    
    if ! command -v qemu-system-aarch64 &> /dev/null; then
        log_error "QEMU is not installed. Install with: brew install qemu"
        exit 1
    fi
    
    local iso_path=$(download_talos_iso)
    
    log_info "Creating controlplane VM..."
    local controlplane_path=$(create_vm_dir "$CONTROLPLANE_NAME")
    local controlplane_disk=$(create_disk_image "$controlplane_path" "$CONTROLPLANE_NAME")
    local controlplane_mac=$(generate_mac "$CONTROLPLANE_NAME")
    local controlplane_script=$(create_vm_script "$controlplane_path" "$CONTROLPLANE_NAME" "$controlplane_mac" "$iso_path" "$controlplane_disk")
    create_service_file "$controlplane_path" "$CONTROLPLANE_NAME" "$controlplane_script"
    
    local manage_script="$VM_DIR/manage-vm.sh"
    cat > "$manage_script" << EOF
#!/bin/bash
VM_DIR="\$(cd "\$(dirname "\$0")" && pwd)"

start_vm() {
    echo "Starting Talos VM..."
    "\$VM_DIR/talos-controlplane/run.sh" &
}

stop_vm() {
    echo "Stopping Talos VM..."
    pkill -f "qemu-system-aarch64.*talos-controlplane"
}

case "\$1" in
    start)
        start_vm
        ;;
    stop)
        stop_vm
        ;;
    status)
        pgrep -f "qemu-system-aarch64.*talos-controlplane" > /dev/null && echo "VM is running" || echo "VM is stopped"
        ;;
    *)
        echo "Usage: \$0 {start|stop|status}"
        exit 1
        ;;
esac
EOF
    chmod +x "$manage_script"
    
    log_info "VM creation complete!"
    log_info "VM files located at: $VM_DIR"
    log_info ""
    log_info "To start VM: $manage_script start"
    log_info "To stop VM: $manage_script stop"
    log_info ""
    log_info "Controlplane MAC: $controlplane_mac"
}

main "$@"