#!/usr/bin/env bash
#
# build-rootfs.sh - Create a minimal Alpine rootfs for Volt testing
#
# This script creates a small, fast-booting root filesystem suitable
# for microVM testing. Uses Alpine Linux for its minimal footprint.
#
# Requirements:
#   - curl, tar
#   - e2fsprogs (mkfs.ext4) or squashfs-tools (mksquashfs)
#   - Optional: sudo (for proper permissions)
#
# Output: images/alpine-rootfs.ext4 (or .squashfs)

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
BUILD_DIR="${PROJECT_DIR}/.build/rootfs"
OUTPUT_DIR="${PROJECT_DIR}/images"

# Alpine version
ALPINE_VERSION="${ALPINE_VERSION:-3.19}"
ALPINE_RELEASE="${ALPINE_RELEASE:-3.19.1}"
ALPINE_ARCH="x86_64"
ALPINE_URL="https://dl-cdn.alpinelinux.org/alpine/v${ALPINE_VERSION}/releases/${ALPINE_ARCH}/alpine-minirootfs-${ALPINE_RELEASE}-${ALPINE_ARCH}.tar.gz"

# Image settings
IMAGE_FORMAT="${IMAGE_FORMAT:-ext4}"  # ext4 or squashfs
IMAGE_SIZE_MB="${IMAGE_SIZE_MB:-64}"  # Size for ext4 images
IMAGE_NAME="alpine-rootfs"

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

log() { echo -e "${GREEN}[+]${NC} $*"; }
warn() { echo -e "${YELLOW}[!]${NC} $*"; }
error() { echo -e "${RED}[✗]${NC} $*"; exit 1; }

check_dependencies() {
    log "Checking dependencies..."
    
    local deps=(curl tar)
    
    case "$IMAGE_FORMAT" in
        ext4) deps+=(mkfs.ext4) ;;
        squashfs) deps+=(mksquashfs) ;;
        *) error "Unknown format: $IMAGE_FORMAT" ;;
    esac
    
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &>/dev/null; then
            error "Missing dependency: $dep"
        fi
    done
}

download_alpine() {
    log "Downloading Alpine minirootfs ${ALPINE_RELEASE}..."
    
    mkdir -p "$BUILD_DIR"
    
    local tarball="${BUILD_DIR}/alpine-minirootfs.tar.gz"
    if [[ ! -f "$tarball" ]]; then
        curl -L -o "$tarball" "$ALPINE_URL"
    else
        log "Using cached download"
    fi
}

extract_rootfs() {
    log "Extracting rootfs..."
    
    local rootfs="${BUILD_DIR}/rootfs"
    rm -rf "$rootfs"
    mkdir -p "$rootfs"
    
    # Extract (needs root for proper permissions, but works without)
    if [[ $EUID -eq 0 ]]; then
        tar xzf "${BUILD_DIR}/alpine-minirootfs.tar.gz" -C "$rootfs"
    else
        # Fakeroot alternative or just extract
        tar xzf "${BUILD_DIR}/alpine-minirootfs.tar.gz" -C "$rootfs" 2>/dev/null || \
            tar xzf "${BUILD_DIR}/alpine-minirootfs.tar.gz" -C "$rootfs" --no-same-owner
        warn "Extracted without root - some permissions may be incorrect"
    fi
}

customize_rootfs() {
    log "Customizing rootfs for microVM boot..."
    
    local rootfs="${BUILD_DIR}/rootfs"
    
    # Create init script for fast boot
    cat > "${rootfs}/init" << 'INIT'
#!/bin/sh
# Volt microVM init

# Mount essential filesystems
mount -t proc proc /proc
mount -t sysfs sys /sys
mount -t devtmpfs dev /dev

# Set hostname
hostname volt-vmm-vm

# Print boot message
echo ""
echo "======================================"
echo "  Volt microVM booted!"
echo "  Alpine Linux $(cat /etc/alpine-release)"
echo "======================================"
echo ""

# Show boot time if available
if [ -f /proc/uptime ]; then
    uptime=$(cut -d' ' -f1 /proc/uptime)
    echo "Boot time: ${uptime}s"
fi

# Start shell
exec /bin/sh
INIT
    chmod +x "${rootfs}/init"
    
    # Create minimal inittab
    cat > "${rootfs}/etc/inittab" << 'EOF'
::sysinit:/etc/init.d/rcS
::respawn:-/bin/sh
ttyS0::respawn:/sbin/getty -L ttyS0 115200 vt100
::shutdown:/bin/umount -a -r
EOF
    
    # Configure serial console
    mkdir -p "${rootfs}/etc/init.d"
    cat > "${rootfs}/etc/init.d/rcS" << 'EOF'
#!/bin/sh
mount -t proc proc /proc
mount -t sysfs sys /sys
mount -t devtmpfs dev /dev
hostname volt-vmm-vm
EOF
    chmod +x "${rootfs}/etc/init.d/rcS"
    
    # Set up basic networking config
    mkdir -p "${rootfs}/etc/network"
    cat > "${rootfs}/etc/network/interfaces" << 'EOF'
auto lo
iface lo inet loopback

auto eth0
iface eth0 inet dhcp
EOF
    
    # Disable unnecessary services
    rm -f "${rootfs}/etc/init.d/hwclock"
    rm -f "${rootfs}/etc/init.d/hwdrivers"
    
    # Create fstab
    cat > "${rootfs}/etc/fstab" << 'EOF'
/dev/vda / ext4 defaults,noatime 0 1
proc /proc proc defaults 0 0
sys /sys sysfs defaults 0 0
devpts /dev/pts devpts defaults 0 0
EOF
    
    log "Rootfs customized for fast boot"
}

create_ext4_image() {
    log "Creating ext4 image (${IMAGE_SIZE_MB}MB)..."
    
    mkdir -p "$OUTPUT_DIR"
    local image="${OUTPUT_DIR}/${IMAGE_NAME}.ext4"
    local rootfs="${BUILD_DIR}/rootfs"
    
    # Create sparse file
    dd if=/dev/zero of="$image" bs=1M count=0 seek="$IMAGE_SIZE_MB" 2>/dev/null
    
    # Format
    mkfs.ext4 -F -L rootfs -O ^metadata_csum "$image" >/dev/null
    
    # Mount and copy (requires root)
    if [[ $EUID -eq 0 ]]; then
        local mnt="${BUILD_DIR}/mnt"
        mkdir -p "$mnt"
        mount -o loop "$image" "$mnt"
        cp -a "${rootfs}/." "$mnt/"
        umount "$mnt"
    else
        # Use debugfs to copy files (limited but works without root)
        warn "Creating image without root - using alternative method"
        
        # Create a tar and extract into image using e2tools or fuse
        if command -v e2cp &>/dev/null; then
            # Use e2tools
            find "$rootfs" -type f | while read -r file; do
                local dest="${file#$rootfs}"
                e2cp "$file" "$image:$dest" 2>/dev/null || true
            done
        else
            warn "e2fsprogs-extra not available - image will be empty"
            warn "Install e2fsprogs-extra or run as root for full rootfs"
        fi
    fi
    
    echo "$image"
}

create_squashfs_image() {
    log "Creating squashfs image..."
    
    mkdir -p "$OUTPUT_DIR"
    local image="${OUTPUT_DIR}/${IMAGE_NAME}.squashfs"
    local rootfs="${BUILD_DIR}/rootfs"
    
    mksquashfs "$rootfs" "$image" \
        -comp zstd \
        -Xcompression-level 19 \
        -noappend \
        -quiet
    
    echo "$image"
}

create_image() {
    local image
    
    case "$IMAGE_FORMAT" in
        ext4) image=$(create_ext4_image) ;;
        squashfs) image=$(create_squashfs_image) ;;
    esac
    
    echo "$image"
}

show_stats() {
    local image="$1"
    
    log "Rootfs image created successfully!"
    echo ""
    echo "  Path: $image"
    echo "  Size: $(du -h "$image" | cut -f1)"
    echo "  Format: $IMAGE_FORMAT"
    echo "  Base: Alpine Linux ${ALPINE_RELEASE}"
    echo ""
    echo "To use with Volt:"
    echo "  volt-vmm --kernel kernels/vmlinux --rootfs $image"
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --format)
            IMAGE_FORMAT="$2"
            shift 2
            ;;
        --size)
            IMAGE_SIZE_MB="$2"
            shift 2
            ;;
        --help)
            echo "Usage: $0 [--format ext4|squashfs] [--size MB]"
            exit 0
            ;;
        *)
            error "Unknown option: $1"
            ;;
    esac
done

# Main
main() {
    log "Building Volt test rootfs"
    echo ""
    
    check_dependencies
    download_alpine
    extract_rootfs
    customize_rootfs
    
    local image
    image=$(create_image)
    
    show_stats "$image"
}

main