#!/bin/bash
#
# Volt Platform - Image Builder
# Creates TinyVol images from definitions
#

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
CONFIG_DIR="$PROJECT_DIR/configs/images"
OUTPUT_DIR="${OUTPUT_DIR:-/var/lib/volt/images}"
CACHE_DIR="${CACHE_DIR:-/var/cache/volt/packages}"

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

log() { echo -e "${GREEN}[volt]${NC} $1"; }
info() { echo -e "${BLUE}[volt]${NC} $1"; }
warn() { echo -e "${YELLOW}[volt]${NC} $1"; }
error() { echo -e "${RED}[volt]${NC} $1" >&2; }

# Base packages for each userland type
declare -A USERLAND_PACKAGES=(
    ["musl-minimal"]="musl busybox"
    ["glibc-standard"]="glibc bash coreutils util-linux systemd"
    ["busybox-tiny"]="busybox-static"
)

# Image definitions
declare -A IMAGES=(
    ["volt/server"]="server.yaml"
    ["volt/server-db-postgres"]="server-db-postgres.yaml"
    ["volt/dev"]="dev.yaml"
    ["volt/desktop-minimal"]="desktop-minimal.yaml"
    ["volt/desktop-productivity"]="desktop-productivity.yaml"
    ["volt/edge"]="edge.yaml"
    ["volt/k8s-node"]="k8s-node.yaml"
)

build_rootfs() {
    local image_name="$1"
    local config_file="$2"
    local rootfs_dir="$3"
    
    log "Building rootfs for: $image_name"
    
    # Create directory structure
    mkdir -p "$rootfs_dir"/{bin,sbin,usr/{bin,sbin,lib},lib,lib64,etc,var,tmp,proc,sys,dev,run,home,root,app}
    
    # Parse YAML config (simplified - in production use proper YAML parser)
    local userland=$(grep "userland:" "$config_file" 2>/dev/null | awk '{print $2}' || echo "musl-minimal")
    
    info "  Userland: $userland"
    
    # Install base userland
    case "$userland" in
        musl-minimal)
            install_musl_minimal "$rootfs_dir"
            ;;
        glibc-standard)
            install_glibc_standard "$rootfs_dir"
            ;;
        busybox-tiny)
            install_busybox_tiny "$rootfs_dir"
            ;;
        *)
            warn "Unknown userland: $userland, using musl-minimal"
            install_musl_minimal "$rootfs_dir"
            ;;
    esac
    
    # Create essential files
    create_essential_files "$rootfs_dir"
    
    # Set permissions
    chmod 1777 "$rootfs_dir/tmp"
    chmod 755 "$rootfs_dir"
}

install_musl_minimal() {
    local rootfs="$1"
    
    info "  Installing musl-minimal userland..."
    
    # Download and install BusyBox static binary
    local busybox_url="https://busybox.net/downloads/binaries/1.35.0-x86_64-linux-musl/busybox"
    local busybox_path="$rootfs/bin/busybox"
    
    if [[ ! -f "$CACHE_DIR/busybox" ]]; then
        mkdir -p "$CACHE_DIR"
        curl -fSL -o "$CACHE_DIR/busybox" "$busybox_url" || {
            # Fallback: create minimal shell script
            warn "Could not download busybox, creating minimal placeholder"
            cat > "$busybox_path" << 'BUSYBOX'
#!/bin/sh
echo "Volt minimal shell"
exec /bin/sh "$@"
BUSYBOX
            chmod +x "$busybox_path"
            return 0
        }
    fi
    
    cp "$CACHE_DIR/busybox" "$busybox_path"
    chmod +x "$busybox_path"
    
    # Create symlinks for common utilities
    local utils="sh ash ls cat cp mv rm mkdir rmdir ln echo pwd env grep sed awk head tail sort uniq wc cut tr sleep date hostname uname id whoami ps kill"
    for util in $utils; do
        ln -sf busybox "$rootfs/bin/$util"
    done
    
    # Create sbin links
    local sbin_utils="init halt reboot poweroff mount umount ifconfig route"
    for util in $sbin_utils; do
        ln -sf ../bin/busybox "$rootfs/sbin/$util"
    done
}

install_glibc_standard() {
    local rootfs="$1"
    
    info "  Installing glibc-standard userland..."
    
    # For now, use Alpine as a base (it's actually musl but good enough for development)
    # In production, this would pull from ArmoredGateHub registry
    
    # Create minimal glibc-like structure
    install_musl_minimal "$rootfs"
    
    # Add bash if available
    if command -v bash &>/dev/null; then
        cp "$(command -v bash)" "$rootfs/bin/bash" 2>/dev/null || true
    fi
    
    # Copy essential libraries from host (for development only)
    # In production, these come from TinyVol images
    for lib in /lib/x86_64-linux-gnu/libc.so.6 /lib/x86_64-linux-gnu/libm.so.6 /lib/x86_64-linux-gnu/libdl.so.2 /lib/x86_64-linux-gnu/libpthread.so.0; do
        if [[ -f "$lib" ]]; then
            mkdir -p "$rootfs/lib/x86_64-linux-gnu"
            cp "$lib" "$rootfs/lib/x86_64-linux-gnu/" 2>/dev/null || true
        fi
    done
    
    # Copy ld-linux
    if [[ -f /lib64/ld-linux-x86-64.so.2 ]]; then
        mkdir -p "$rootfs/lib64"
        cp /lib64/ld-linux-x86-64.so.2 "$rootfs/lib64/" 2>/dev/null || true
    fi
}

install_busybox_tiny() {
    local rootfs="$1"
    
    info "  Installing busybox-tiny userland..."
    
    # Absolute minimal - just busybox
    install_musl_minimal "$rootfs"
    
    # Remove non-essential symlinks
    rm -f "$rootfs/bin/awk" "$rootfs/bin/sed" "$rootfs/bin/grep"
}

create_essential_files() {
    local rootfs="$1"
    
    # /etc/passwd
    cat > "$rootfs/etc/passwd" << 'EOF'
root:x:0:0:root:/root:/bin/sh
nobody:x:65534:65534:nobody:/nonexistent:/usr/sbin/nologin
volt:x:1000:1000:Volt User:/home/volt:/bin/sh
EOF

    # /etc/group
    cat > "$rootfs/etc/group" << 'EOF'
root:x:0:
nobody:x:65534:
volt:x:1000:
EOF

    # /etc/shadow (empty passwords - VMs use keys)
    cat > "$rootfs/etc/shadow" << 'EOF'
root:*:19000:0:99999:7:::
nobody:*:19000:0:99999:7:::
volt:*:19000:0:99999:7:::
EOF
    chmod 640 "$rootfs/etc/shadow"

    # /etc/hosts
    cat > "$rootfs/etc/hosts" << 'EOF'
127.0.0.1   localhost
::1         localhost ip6-localhost ip6-loopback
EOF

    # /etc/hostname
    echo "volt" > "$rootfs/etc/hostname"

    # /etc/resolv.conf
    cat > "$rootfs/etc/resolv.conf" << 'EOF'
nameserver 8.8.8.8
nameserver 8.8.4.4
EOF

    # /etc/nsswitch.conf
    cat > "$rootfs/etc/nsswitch.conf" << 'EOF'
passwd:     files
group:      files
shadow:     files
hosts:      files dns
networks:   files
protocols:  files
services:   files
EOF

    # /etc/os-release
    cat > "$rootfs/etc/os-release" << 'EOF'
NAME="Volt Platform"
VERSION="1.0"
ID=volt
ID_LIKE=alpine
VERSION_ID=1.0
PRETTY_NAME="Volt Platform VM"
HOME_URL="https://voltvisor.io"
EOF

    # Init script
    cat > "$rootfs/sbin/init" << 'INIT'
#!/bin/sh
# Volt Init

# Mount essential filesystems
mount -t proc proc /proc
mount -t sysfs sysfs /sys
mount -t devtmpfs devtmpfs /dev
mkdir -p /dev/pts /dev/shm
mount -t devpts devpts /dev/pts
mount -t tmpfs tmpfs /dev/shm
mount -t tmpfs tmpfs /tmp
mount -t tmpfs tmpfs /run

# Set hostname
hostname -F /etc/hostname

# Network (if configured)
if [ -f /etc/network/interfaces ]; then
    ifconfig lo up
fi

# Run init scripts
for script in /etc/init.d/S*; do
    [ -x "$script" ] && "$script" start
done

# Start shell or configured service
if [ -f /etc/volt/service ]; then
    exec $(cat /etc/volt/service)
else
    exec /bin/sh
fi
INIT
    chmod +x "$rootfs/sbin/init"
    
    # Create init.d directory
    mkdir -p "$rootfs/etc/init.d"
    mkdir -p "$rootfs/etc/volt"
}

create_tinyvol() {
    local image_name="$1"
    local rootfs_dir="$2"
    local output_path="$3"
    
    log "Creating TinyVol: $output_path"
    
    # Create squashfs image (TinyVol format)
    # In production, this would use the actual TinyVol format
    if command -v mksquashfs &>/dev/null; then
        mksquashfs "$rootfs_dir" "$output_path" \
            -comp zstd \
            -Xcompression-level 19 \
            -all-root \
            -noappend \
            -no-progress
    else
        # Fallback: create tar archive
        warn "mksquashfs not found, creating tar archive"
        tar -czf "$output_path" -C "$rootfs_dir" .
    fi
    
    local size=$(du -h "$output_path" | cut -f1)
    info "  Image size: $size"
}

generate_sbom() {
    local image_name="$1"
    local rootfs_dir="$2"
    local output_path="$3"
    
    log "Generating SBOM for: $image_name"
    
    cat > "$output_path" << EOF
{
    "bomFormat": "CycloneDX",
    "specVersion": "1.5",
    "version": 1,
    "metadata": {
        "timestamp": "$(date -Iseconds)",
        "component": {
            "type": "operating-system",
            "name": "$image_name",
            "version": "1.0"
        }
    },
    "components": [
        {
            "type": "application",
            "name": "busybox",
            "version": "1.35.0"
        }
    ]
}
EOF
}

sign_image() {
    local image_name="$1"
    local image_path="$2"
    
    log "Signing image: $image_name"
    
    # Generate checksums
    sha256sum "$image_path" > "${image_path}.sha256"
    
    # TODO: Integrate with ArmoredForge signing
    # armored-forge sign "$image_path" --key volt-image-key
}

build_image() {
    local image_name="$1"
    local config_file="$CONFIG_DIR/$2"
    
    log "=========================================="
    log "Building image: $image_name"
    log "=========================================="
    
    if [[ ! -f "$config_file" ]]; then
        warn "Config file not found: $config_file"
        # Create default config
        config_file="$CONFIG_DIR/server.yaml"
    fi
    
    local safe_name=$(echo "$image_name" | tr '/' '_')
    local work_dir="$OUTPUT_DIR/.build/$safe_name"
    local rootfs_dir="$work_dir/rootfs"
    local image_path="$OUTPUT_DIR/$safe_name.tinyvol"
    
    # Clean and create work directory
    rm -rf "$work_dir"
    mkdir -p "$work_dir" "$rootfs_dir"
    
    # Build rootfs
    build_rootfs "$image_name" "$config_file" "$rootfs_dir"
    
    # Create TinyVol image
    create_tinyvol "$image_name" "$rootfs_dir" "$image_path"
    
    # Generate SBOM
    generate_sbom "$image_name" "$rootfs_dir" "${image_path}.sbom.json"
    
    # Sign image
    sign_image "$image_name" "$image_path"
    
    # Create image metadata
    cat > "${image_path}.json" << EOF
{
    "name": "$image_name",
    "version": "1.0",
    "created": "$(date -Iseconds)",
    "size": "$(du -h "$image_path" | cut -f1)",
    "sha256": "$(sha256sum "$image_path" | cut -d' ' -f1)",
    "sbom": "${image_path}.sbom.json"
}
EOF
    
    # Cleanup work directory
    rm -rf "$work_dir"
    
    log "Image built: $image_path"
}

main() {
    log "Volt Platform Image Builder"
    log "=============================="
    
    mkdir -p "$OUTPUT_DIR" "$CACHE_DIR"
    
    # Build all defined images
    for image_name in "${!IMAGES[@]}"; do
        build_image "$image_name" "${IMAGES[$image_name]}"
    done
    
    # If no images defined, build defaults
    if [[ ${#IMAGES[@]} -eq 0 ]]; then
        build_image "volt/server" "server.yaml"
        build_image "volt/desktop-productivity" "desktop-productivity.yaml"
    fi
    
    log ""
    log "Build complete!"
    log "Images installed to: $OUTPUT_DIR"
    ls -la "$OUTPUT_DIR"/*.tinyvol 2>/dev/null || ls -la "$OUTPUT_DIR"
}

# Run if executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi