#!/bin/bash

# Alpine Linux Boot Script with Cloud-Hypervisor and VirtioFS
# This script downloads Alpine Linux miniroot, kernel, and initrd,
# then boots using cloud-hypervisor with virtiofs

set -euo pipefail

# Configuration
ALPINE_VERSION="3.22"
ALPINE_ARCH="x86_64"
ALPINE_MIRROR="https://dl-cdn.alpinelinux.org/alpine"
WORK_DIR="$(pwd)/alpine-boot"
MINIROOT_DIR="$WORK_DIR/miniroot"
KERNEL_DIR="$WORK_DIR/kernel"

# Default hostname
DEFAULT_HOSTNAME="alpine-$ALPINE_VERSION"
HOSTNAME="$DEFAULT_HOSTNAME"

# Root password (use --root-password to override)
ROOT_PASSWORD="root"

# URLs
ALPINE_BASE_URL="$ALPINE_MIRROR/v$ALPINE_VERSION/releases/$ALPINE_ARCH"
MINIROOT_URL="$ALPINE_BASE_URL/alpine-minirootfs-$ALPINE_VERSION.0-$ALPINE_ARCH.tar.gz"
NETBOOT_URL="$ALPINE_BASE_URL/netboot-$ALPINE_VERSION.0"
VMLINUZ_URL="$NETBOOT_URL/vmlinuz-virt"
INITRAMFS_URL="$NETBOOT_URL/initramfs-virt"

# Cloud-hypervisor configuration
VM_MEMORY="1024M"
VM_CPUS="2"
KERNEL_CMDLINE="console=ttyS0 rootfstype=virtiofs root=/dev/root debug rw init=/sbin/init debug_init"

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

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

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

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

# Find virtiofsd location
find_virtiofsd() {
    # Check common locations for virtiofsd
    local virtiofsd_paths=(
        "virtiofsd"  # In PATH
        "/usr/lib/virtiofsd"  # Arch Linux
        "/usr/libexec/virtiofsd"  # Some distributions
        "/usr/bin/virtiofsd"  # Standard location
    )
    
    for path in "${virtiofsd_paths[@]}"; do
        if [[ "$path" == "virtiofsd" ]]; then
            # Check if it's in PATH
            if command -v "$path" >/dev/null 2>&1; then
                echo "$path"
                return 0
            fi
        else
            # Check if file exists for absolute paths
            if [[ -f "$path" && -x "$path" ]]; then
                echo "$path"
                return 0
            fi
        fi
    done
    
    return 1
}

# Check dependencies
check_dependencies() {
    log "Checking dependencies..."
    
    local missing_deps=()
    
    command -v curl >/dev/null 2>&1 || missing_deps+=("curl")
    command -v tar >/dev/null 2>&1 || missing_deps+=("tar")
    command -v cloud-hypervisor >/dev/null 2>&1 || missing_deps+=("cloud-hypervisor")
    
    # Check for virtiofsd
    if ! VIRTIOFSD_CMD=$(find_virtiofsd); then
        missing_deps+=("virtiofsd")
    else
        log "Found virtiofsd at: $VIRTIOFSD_CMD"
    fi
    
    if [ ${#missing_deps[@]} -ne 0 ]; then
        error "Missing dependencies: ${missing_deps[*]}"
    fi
    
    log "All dependencies found."
}

# Create working directories
setup_directories() {
    log "Setting up directories..."
    mkdir -p "$WORK_DIR" "$MINIROOT_DIR" "$KERNEL_DIR"
}

# Download and extract miniroot
download_miniroot() {
    log "Downloading Alpine Linux miniroot..."
    
    local miniroot_file="$WORK_DIR/alpine-minirootfs.tar.gz"
    
    if [ ! -f "$miniroot_file" ]; then
        curl -L -o "$miniroot_file" "$MINIROOT_URL" || error "Failed to download miniroot"
    else
        log "Miniroot already downloaded, skipping..."
    fi
    
    log "Extracting miniroot..."
    if [ ! -d "$MINIROOT_DIR/bin" ]; then
        tar -xzf "$miniroot_file" -C "$MINIROOT_DIR" || error "Failed to extract miniroot"
    else
        log "Miniroot already extracted, skipping..."
    fi
}

# Download kernel and initramfs
download_kernel() {
    log "Downloading Alpine Linux kernel and initramfs..."
    
    local vmlinuz_file="$KERNEL_DIR/vmlinuz"
    local initramfs_file="$KERNEL_DIR/initramfs"
    
    # Download vmlinuz
    if [ ! -f "$vmlinuz_file" ]; then
        log "Downloading vmlinuz..."
        curl -L -o "$vmlinuz_file" "$VMLINUZ_URL" || error "Failed to download vmlinuz"
    else
        log "vmlinuz already downloaded, skipping..."
    fi
    
    # Download initramfs
    if [ ! -f "$initramfs_file" ]; then
        log "Downloading initramfs..."
        curl -L -o "$initramfs_file" "$INITRAMFS_URL" || error "Failed to download initramfs"
    else
        log "initramfs already downloaded, skipping..."
    fi
    
    # Verify files
    if [ ! -f "$vmlinuz_file" ] || [ ! -s "$vmlinuz_file" ]; then
        error "vmlinuz file is missing or empty"
    fi
    
    if [ ! -f "$initramfs_file" ] || [ ! -s "$initramfs_file" ]; then
        error "initramfs file is missing or empty"
    fi
}

# Customize miniroot with chroot
customize_miniroot() {
    log "Customizing miniroot with additional packages and configuration..."
    
    # Set up chroot environment
    log "Setting up chroot environment..."
    sudo mount --bind /proc "$MINIROOT_DIR/proc" || error "Failed to bind mount /proc"
    sudo mount --bind /sys "$MINIROOT_DIR/sys" || error "Failed to bind mount /sys"
    sudo mount --bind /dev "$MINIROOT_DIR/dev" || error "Failed to bind mount /dev"
    
    # Cleanup function for chroot environment
    cleanup_chroot() {
        log "Cleaning up chroot environment..."
        sudo umount "$MINIROOT_DIR/proc" 2>/dev/null || true
        sudo umount "$MINIROOT_DIR/sys" 2>/dev/null || true
        sudo umount "$MINIROOT_DIR/dev" 2>/dev/null || true
    }
    trap cleanup_chroot EXIT
    
    # Add hvc0 getty to /etc/inittab
    log "Adding hvc0 getty to /etc/inittab..."
    if ! grep -q "hvc0::respawn" "$MINIROOT_DIR/etc/inittab" 2>/dev/null; then
        sudo chroot "$MINIROOT_DIR" /bin/sh -c "echo 'hvc0::respawn:/sbin/getty -L hvc0 115200 vt100' >> /etc/inittab" || error "Failed to add hvc0 getty"
    else
        log "hvc0 getty already configured, skipping..."
    fi
    
    # Configure nameserver
    log "Configuring nameserver..."
    sudo chroot "$MINIROOT_DIR" /bin/sh -c "echo 'nameserver 1.1.1.1' > /etc/resolv.conf" || error "Failed to configure nameserver"
    
    # Configure hostname (removed from here, will be set separately)
    
    # Update package database and install packages
    log "Updating package database..."
    sudo chroot "$MINIROOT_DIR" /bin/sh -c "apk update" || error "Failed to update package database"
    
    log "Installing openrc, bash, openssh..."
    sudo chroot "$MINIROOT_DIR" /bin/sh -c "apk add openrc bash openssh" || error "Failed to install packages"
    
    # Enable services
    log "Enabling services (sshd, networking)..."
    sudo chroot "$MINIROOT_DIR" /bin/sh -c "rc-update add sshd default" || warn "Failed to enable sshd service"
    sudo chroot "$MINIROOT_DIR" /bin/sh -c "rc-update add networking boot" || warn "Failed to enable networking service"

   # Set root password
   log "Setting root password..."
   warn "Setting default root password. Use --root-password to override. For production, use SSH keys."
   sudo chroot "$MINIROOT_DIR" /bin/sh -c "echo 'root:$ROOT_PASSWORD' | chpasswd" || error "Failed to set root password"
    
    # Enable root SSH login
    log "Enabling root SSH login with password..."
    if [ -f "$MINIROOT_DIR/etc/ssh/sshd_config" ]; then
        sudo sed -i 's/^#PermitRootLogin.*/PermitRootLogin yes/' "$MINIROOT_DIR/etc/ssh/sshd_config" || warn "Could not enable root SSH login"
    else
        warn "sshd_config not found, skipping SSH configuration."
    fi
    
    # Clean up chroot environment
    cleanup_chroot
    trap - EXIT
    
    log "Miniroot customization complete."
}

# Set hostname in miniroot
set_hostname() {
    log "Setting hostname to: $HOSTNAME"
    sudo chroot "$MINIROOT_DIR" /bin/sh -c "echo '$HOSTNAME' > /etc/hostname" || error "Failed to set hostname"
    sudo chroot "$MINIROOT_DIR" /bin/sh -c "echo '127.0.0.1 localhost $HOSTNAME' > /etc/hosts" || error "Failed to configure hosts file"
}

# Prepare miniroot for virtiofs
prepare_miniroot() {
    log "Preparing miniroot for virtiofs..."
    
    # Ensure necessary directories exist
    mkdir -p "$MINIROOT_DIR/dev" "$MINIROOT_DIR/proc" "$MINIROOT_DIR/sys" "$MINIROOT_DIR/tmp"
    
    # Customize miniroot if not already done
    # Use a versioned marker file to ensure re-customization when script logic changes.
    if [ ! -f "$MINIROOT_DIR/.customized_v2" ]; then
        customize_miniroot
        # Mark as customized
        rm -f "$MINIROOT_DIR/.customized" # remove old marker
        touch "$MINIROOT_DIR/.customized_v2"
    else
        log "Miniroot already customized, skipping..."
    fi
    
    # Always set hostname (even if already customized)
    set_hostname
}

# Boot with cloud-hypervisor
boot_vm() {
    log "Booting Alpine Linux with cloud-hypervisor..."
    
    # Clean up stale sockets from previous runs to prevent "Address in use" errors
    log "Cleaning up any stale sockets..."
    rm -f /tmp/ch-virtiofs.sock /tmp/ch-api.sock
    
    local kernel_file="$KERNEL_DIR/vmlinuz"
    local initramfs_file="$KERNEL_DIR/initramfs"
    
    if [ ! -f "$kernel_file" ]; then
        error "Kernel file not found: $kernel_file"
    fi
    
    # Build cloud-hypervisor command
    local ch_cmd=(
        "cloud-hypervisor"
        "--memory" "size=$VM_MEMORY,shared=on"
        "--cpus" "boot=$VM_CPUS"
        "--kernel" "$kernel_file"
        "--cmdline" "$KERNEL_CMDLINE"
        "--fs" "tag=/dev/root,socket=/tmp/ch-virtiofs.sock,num_queues=1"
        "--serial" "tty"
        "--api-socket" "/tmp/ch-api.sock"
    )
    
    # Add initramfs if available
    if [ -f "$initramfs_file" ]; then
        ch_cmd+=("--initramfs" "$initramfs_file")
    fi
    
    log "Starting virtiofsd..."
    # Start virtiofsd in background
    "$VIRTIOFSD_CMD" --socket-path=/tmp/ch-virtiofs.sock --shared-dir="$MINIROOT_DIR" --announce-submounts &
    local virtiofsd_pid=$!
    
    # Give virtiofsd time to start
    sleep 2
    
    log "Starting cloud-hypervisor..."
    log "Command: ${ch_cmd[*]}"
    
    # Cleanup function
    cleanup() {
        log "Cleaning up..."
        kill $virtiofsd_pid 2>/dev/null || true
        rm -f /tmp/ch-virtiofs.sock /tmp/ch-api.sock
    }
    trap cleanup EXIT
    
    # Execute cloud-hypervisor
    "${ch_cmd[@]}" || error "Failed to start cloud-hypervisor"
}

# Main execution
main() {
    log "Alpine Linux Boot Script Starting..."
    
    check_dependencies
    setup_directories
    download_miniroot
    download_kernel
    prepare_miniroot
    boot_vm
}

# Main entry point with argument parsing
handle_args() {
    # Default action
    local action="boot"

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --hostname)
                [[ -z "$2" ]] && error "Missing value for --hostname"
                HOSTNAME="$2"
                shift 2
                ;;
            --root-password)
                [[ -z "$2" ]] && error "Missing value for --root-password"
                ROOT_PASSWORD="$2"
                shift 2
                ;;
            clean|download-only)
                action="$1"
                shift
                ;;
            help|--help|-h)
                action="help"
                shift
                ;;
            *)
                error "Unknown option or command: $1"
                ;;
        esac
    done

    # Execute action
    case "$action" in
        boot)
            log "Using hostname: $HOSTNAME"
            main
            ;;
        clean)
            log "Cleaning up work directory..."
            sudo rm -rf "$WORK_DIR"
            log "Cleanup complete."
            ;;
        download-only)
            log "Download-only mode..."
            check_dependencies
            setup_directories
            download_miniroot
            download_kernel
            log "Downloads complete."
            ;;
        help)
            echo "Usage: $0 [OPTIONS] [COMMAND]"
            echo ""
            echo "OPTIONS:"
            echo "  --hostname HOSTNAME      Set VM hostname (default: $DEFAULT_HOSTNAME)"
            echo "  --root-password PWD      Set root password for the VM (default: 'root')"
            echo ""
            echo "COMMANDS:"
            echo "  clean                    Remove downloaded files and work directory"
            echo "  download-only            Only download files, don't boot"
            echo "  help                     Show this help message"
            echo ""
            echo "Examples:"
            echo "  $0                       Boot with default hostname"
            echo "  $0 --hostname myserver   Boot with hostname 'myserver'"
            echo "  $0 --root-password secret boot with custom password"
            ;;
    esac
}

handle_args "$@"