#!/usr/bin/env bash
# ==============================================================================
# PlasmaCloud NixOS Netboot Image Builder
# ==============================================================================
# This script builds netboot images for bare-metal provisioning of PlasmaCloud.
#
# Usage:
#   ./build-images.sh [--profile PROFILE] [--output-dir DIR] [--help]
#
# Options:
#   --profile PROFILE    Build specific profile (control-plane, worker, all-in-one, all)
#   --output-dir DIR     Output directory for built artifacts (default: ./artifacts)
#   --help              Show this help message
#
# Examples:
#   ./build-images.sh                              # Build all profiles
#   ./build-images.sh --profile control-plane      # Build control plane only
#   ./build-images.sh --profile all                # Build all profiles
#   ./build-images.sh --output-dir /srv/pxe        # Custom output directory
# ==============================================================================

set -euo pipefail

# ==============================================================================
# CONFIGURATION
# ==============================================================================
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
DEFAULT_OUTPUT_DIR="$SCRIPT_DIR/artifacts"
PXE_ASSETS_DIR="$REPO_ROOT/chainfire/baremetal/pxe-server/assets"

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

# ==============================================================================
# FUNCTIONS
# ==============================================================================

# Print colored messages
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

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

# Print banner
print_banner() {
    echo ""
    echo "╔════════════════════════════════════════════════════════════════╗"
    echo "║  PlasmaCloud NixOS Netboot Image Builder                      ║"
    echo "║  Building bare-metal provisioning images                      ║"
    echo "╚════════════════════════════════════════════════════════════════╝"
    echo ""
}

# Print usage
print_usage() {
    cat << EOF
Usage: $0 [OPTIONS]

Build NixOS netboot images for PlasmaCloud bare-metal provisioning.

OPTIONS:
    --profile PROFILE    Build specific profile:
                           - control-plane: All 8 PlasmaCloud services
                           - worker: Compute-focused services (PlasmaVMC, PrismNET)
                           - all-in-one: All services for single-node deployment
                           - all: Build all profiles (default)

    --output-dir DIR     Output directory for artifacts (default: ./artifacts)

    --help              Show this help message

EXAMPLES:
    # Build all profiles
    $0

    # Build control plane only
    $0 --profile control-plane

    # Build to custom output directory
    $0 --output-dir /srv/pxe/images

PROFILES:
    control-plane    - Full control plane with all 8 services
    worker           - Worker node with PlasmaVMC and PrismNET
    all-in-one       - Single-node deployment with all services

OUTPUT:
    The script generates the following artifacts for each profile:
      - bzImage          Linux kernel
      - initrd           Initial ramdisk
      - netboot.ipxe     iPXE boot script

ENVIRONMENT:
    PLASMACLOUD_DEPLOYER_URL     Optional deployer endpoint embedded into generated netboot.ipxe
    PLASMACLOUD_BOOTSTRAP_TOKEN  Optional bootstrap token embedded into generated netboot.ipxe
    PLASMACLOUD_CA_CERT_URL      Optional CA certificate URL embedded into generated netboot.ipxe

EOF
}

# Build a single netboot profile
build_profile() {
    local profile=$1
    local output_dir=$2

    print_info "Building netboot image for profile: $profile"

    # Create profile output directory
    local profile_dir="$output_dir/$profile"
    mkdir -p "$profile_dir"

    # Build the netboot ramdisk
    print_info "  Building initial ramdisk..."
    if ! nix build "$REPO_ROOT#nixosConfigurations.netboot-$profile.config.system.build.netbootRamdisk" \
         --out-link "$profile_dir/initrd-link" 2>&1 | tee "$profile_dir/build.log"; then
        print_error "Failed to build initrd for $profile (see $profile_dir/build.log)"
        return 1
    fi

    # Build the kernel
    print_info "  Building kernel..."
    if ! nix build "$REPO_ROOT#nixosConfigurations.netboot-$profile.config.system.build.kernel" \
         --out-link "$profile_dir/kernel-link" 2>&1 | tee -a "$profile_dir/build.log"; then
        print_error "Failed to build kernel for $profile (see $profile_dir/build.log)"
        return 1
    fi

    # Copy artifacts
    print_info "  Copying artifacts..."
    cp -f "$profile_dir/initrd-link/initrd" "$profile_dir/initrd"
    cp -f "$profile_dir/kernel-link/bzImage" "$profile_dir/bzImage"

    # Resolve init path from the build (avoids hardcoding store paths)
    local init_path="/init"
    if toplevel=$(nix eval --raw "$REPO_ROOT#nixosConfigurations.netboot-$profile.config.system.build.toplevel" 2>/dev/null); then
        if [ -n "$toplevel" ]; then
            init_path="${toplevel}/init"
        fi
    else
        print_warning "Failed to resolve init path for $profile; using /init"
    fi

    local deployer_kernel_args=""
    if [ -n "${PLASMACLOUD_DEPLOYER_URL:-}" ]; then
        deployer_kernel_args+=" plasmacloud.deployer_url=${PLASMACLOUD_DEPLOYER_URL}"
    fi
    if [ -n "${PLASMACLOUD_BOOTSTRAP_TOKEN:-}" ]; then
        deployer_kernel_args+=" plasmacloud.bootstrap_token=${PLASMACLOUD_BOOTSTRAP_TOKEN}"
    fi
    if [ -n "${PLASMACLOUD_CA_CERT_URL:-}" ]; then
        deployer_kernel_args+=" plasmacloud.ca_cert_url=${PLASMACLOUD_CA_CERT_URL}"
    fi

    # Generate iPXE boot script
    print_info "  Generating iPXE boot script..."
    cat > "$profile_dir/netboot.ipxe" << EOF
#!ipxe

# PlasmaCloud Netboot - $profile
# Generated: $(date -u +"%Y-%m-%d %H:%M:%S UTC")

# Set variables
set boot-server \${boot-url}

# Display info
echo Loading PlasmaCloud ($profile profile)...
echo Kernel: bzImage
echo Initrd: initrd
echo

# Load kernel and initrd
kernel \${boot-server}/$profile/bzImage init=${init_path} console=ttyS0,115200 console=tty0 loglevel=4${deployer_kernel_args}
initrd \${boot-server}/$profile/initrd

# Boot
boot
EOF

    # Calculate sizes
    local kernel_size=$(du -h "$profile_dir/bzImage" | cut -f1)
    local initrd_size=$(du -h "$profile_dir/initrd" | cut -f1)
    local total_size=$(du -sh "$profile_dir" | cut -f1)

    # Print summary
    print_success "Profile $profile built successfully!"
    print_info "  Kernel:     $kernel_size"
    print_info "  Initrd:     $initrd_size"
    print_info "  Total:      $total_size"
    print_info "  Location:   $profile_dir"
    echo ""
}

# Copy artifacts to PXE server assets directory
copy_to_pxe_server() {
    local output_dir=$1

    if [ ! -d "$PXE_ASSETS_DIR" ]; then
        print_warning "PXE assets directory not found: $PXE_ASSETS_DIR"
        print_warning "Skipping copy to PXE server"
        return 0
    fi

    print_info "Copying artifacts to PXE server: $PXE_ASSETS_DIR"

    for profile in control-plane worker all-in-one; do
        local profile_dir="$output_dir/$profile"
        if [ -d "$profile_dir" ]; then
            local pxe_profile_dir="$PXE_ASSETS_DIR/nixos/$profile"
            mkdir -p "$pxe_profile_dir"

            cp -f "$profile_dir/bzImage" "$pxe_profile_dir/"
            cp -f "$profile_dir/initrd" "$pxe_profile_dir/"
            cp -f "$profile_dir/netboot.ipxe" "$pxe_profile_dir/"

            # Create symlinks for convenience
            ln -sf "$profile/bzImage" "$PXE_ASSETS_DIR/nixos/bzImage-$profile"
            ln -sf "$profile/initrd" "$PXE_ASSETS_DIR/nixos/initrd-$profile"

            print_success "  Copied $profile to PXE server"
        fi
    done
}

# Verify build outputs
verify_outputs() {
    local output_dir=$1
    local profile=$2

    local profile_dir="$output_dir/$profile"
    local errors=0

    # Check for required files
    if [ ! -f "$profile_dir/bzImage" ]; then
        print_error "Missing bzImage for $profile"
        ((errors++))
    fi

    if [ ! -f "$profile_dir/initrd" ]; then
        print_error "Missing initrd for $profile"
        ((errors++))
    fi

    if [ ! -f "$profile_dir/netboot.ipxe" ]; then
        print_error "Missing netboot.ipxe for $profile"
        ((errors++))
    fi

    # Check file sizes (should be reasonable)
    if [ -f "$profile_dir/bzImage" ]; then
        local kernel_size=$(stat -c%s "$profile_dir/bzImage")
        if [ "$kernel_size" -lt 1000000 ]; then  # Less than 1MB is suspicious
            print_warning "Kernel size seems too small: $kernel_size bytes"
            ((errors++))
        fi
    fi

    if [ -f "$profile_dir/initrd" ]; then
        local initrd_size=$(stat -c%s "$profile_dir/initrd")
        if [ "$initrd_size" -lt 10000000 ]; then  # Less than 10MB is suspicious
            print_warning "Initrd size seems too small: $initrd_size bytes"
            ((errors++))
        fi
    fi

    return $errors
}

# Print final summary
print_summary() {
    local output_dir=$1
    local profiles=("$@")
    shift  # Remove first argument (output_dir)

    echo ""
    echo "╔════════════════════════════════════════════════════════════════╗"
    echo "║  Build Summary                                                 ║"
    echo "╚════════════════════════════════════════════════════════════════╝"
    echo ""

    for profile in "${profiles[@]}"; do
        if [ "$profile" == "$output_dir" ]; then
            continue
        fi

        local profile_dir="$output_dir/$profile"
        if [ -d "$profile_dir" ]; then
            echo "Profile: $profile"
            echo "  Location: $profile_dir"
            if [ -f "$profile_dir/bzImage" ]; then
                echo "  Kernel:   $(du -h "$profile_dir/bzImage" | cut -f1)"
            fi
            if [ -f "$profile_dir/initrd" ]; then
                echo "  Initrd:   $(du -h "$profile_dir/initrd" | cut -f1)"
            fi
            echo ""
        fi
    done

    echo "Next Steps:"
    echo "  1. Deploy images to PXE server (if not done automatically)"
    echo "  2. Configure DHCP/iPXE boot infrastructure"
    echo "  3. Boot target machines via PXE"
    echo "  4. Use nixos-anywhere for installation"
    echo ""
    echo "For more information, see:"
    echo "  - baremetal/image-builder/README.md"
    echo "  - docs/por/T032-baremetal-provisioning/design.md"
    echo ""
}

# ==============================================================================
# MAIN
# ==============================================================================

main() {
    local profile="all"
    local output_dir="$DEFAULT_OUTPUT_DIR"

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --profile)
                profile="$2"
                shift 2
                ;;
            --output-dir)
                output_dir="$2"
                shift 2
                ;;
            --help)
                print_usage
                exit 0
                ;;
            *)
                print_error "Unknown option: $1"
                print_usage
                exit 1
                ;;
        esac
    done

    # Validate profile
    if [[ ! "$profile" =~ ^(control-plane|worker|all-in-one|all)$ ]]; then
        print_error "Invalid profile: $profile"
        print_usage
        exit 1
    fi

    print_banner

    # Create output directory
    mkdir -p "$output_dir"

    # Build profiles
    local profiles_to_build=()
    if [ "$profile" == "all" ]; then
        profiles_to_build=("control-plane" "worker" "all-in-one")
    else
        profiles_to_build=("$profile")
    fi

    local build_errors=0

    for p in "${profiles_to_build[@]}"; do
        if ! build_profile "$p" "$output_dir"; then
            ((build_errors++))
        fi
    done

    # Verify outputs
    print_info "Verifying build outputs..."
    local verify_errors=0
    for p in "${profiles_to_build[@]}"; do
        if ! verify_outputs "$output_dir" "$p"; then
            ((verify_errors++))
        fi
    done

    # Copy to PXE server if available
    copy_to_pxe_server "$output_dir"

    # Print summary
    print_summary "$output_dir" "${profiles_to_build[@]}"

    # Exit with error if any builds failed
    if [ $build_errors -gt 0 ]; then
        print_error "Build completed with $build_errors error(s)"
        exit 1
    fi

    if [ $verify_errors -gt 0 ]; then
        print_warning "Build completed with $verify_errors warning(s)"
    fi

    print_success "All builds completed successfully!"
}

# Run main function
main "$@"