#!/usr/bin/env bash
###############################################################################
# PXE Boot Server Setup Script
#
# This script prepares a PXE boot server for Centra Cloud bare-metal
# provisioning. It performs the following tasks:
#
# 1. Creates directory structure for boot assets
# 2. Downloads iPXE bootloaders (or provides build instructions)
# 3. Copies configuration files to appropriate locations
# 4. Validates configuration files
# 5. Tests DHCP/TFTP/HTTP services
#
# Usage:
#   sudo ./setup.sh [options]
#
# Options:
#   --install         Install and configure services
#   --download        Download iPXE bootloaders
#   --build-ipxe      Build iPXE from source (recommended for production)
#   --validate        Validate configuration files
#   --test            Test services (DHCP, TFTP, HTTP)
#   --help            Show this help message
#
# Example:
#   sudo ./setup.sh --install --download --validate
###############################################################################

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BOOT_ASSETS_DIR="/var/lib/pxe-boot"
IPXE_DIR="${BOOT_ASSETS_DIR}/ipxe"
NIXOS_DIR="${BOOT_ASSETS_DIR}/nixos"
TFTP_DIR="/var/lib/tftpboot"

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

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $*"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $*"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $*"
}

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

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root (use sudo)"
        exit 1
    fi
}

# Display help
show_help() {
    cat << EOF
PXE Boot Server Setup Script

Usage: sudo $0 [options]

Options:
    --install         Install and configure services
    --download        Download iPXE bootloaders from boot.ipxe.org
    --build-ipxe      Build iPXE from source (recommended for production)
    --validate        Validate configuration files
    --test            Test services (DHCP, TFTP, HTTP)
    --clean           Clean up boot assets directory
    --help            Show this help message

Examples:
    # Full setup with pre-built bootloaders
    sudo $0 --install --download --validate

    # Build iPXE from source (more secure, customizable)
    sudo $0 --install --build-ipxe --validate

    # Validate configuration only
    sudo $0 --validate

    # Test services
    sudo $0 --test

For more information, see README.md
EOF
}

# Create directory structure
create_directories() {
    log_info "Creating directory structure..."

    mkdir -p "${IPXE_DIR}"
    mkdir -p "${NIXOS_DIR}"
    mkdir -p "${TFTP_DIR}"
    mkdir -p /var/log/dhcpd
    mkdir -p /var/log/nginx

    # Set permissions
    chown -R nginx:nginx "${BOOT_ASSETS_DIR}" 2>/dev/null || \
        log_warning "nginx user not found, skipping chown (install nginx first)"

    chmod -R 755 "${BOOT_ASSETS_DIR}"

    log_success "Directory structure created at ${BOOT_ASSETS_DIR}"
}

# Download iPXE bootloaders
download_ipxe() {
    log_info "Downloading iPXE bootloaders from boot.ipxe.org..."

    # URLs for iPXE bootloaders
    IPXE_BASE_URL="https://boot.ipxe.org"

    # Download BIOS bootloader (undionly.kpxe)
    if [[ ! -f "${IPXE_DIR}/undionly.kpxe" ]]; then
        log_info "Downloading undionly.kpxe (BIOS bootloader)..."
        curl -L -o "${IPXE_DIR}/undionly.kpxe" "${IPXE_BASE_URL}/undionly.kpxe" || {
            log_error "Failed to download undionly.kpxe"
            return 1
        }
        log_success "Downloaded undionly.kpxe ($(du -h "${IPXE_DIR}/undionly.kpxe" | cut -f1))"
    else
        log_info "undionly.kpxe already exists, skipping download"
    fi

    # Download UEFI bootloader (ipxe.efi)
    if [[ ! -f "${IPXE_DIR}/ipxe.efi" ]]; then
        log_info "Downloading ipxe.efi (UEFI x86-64 bootloader)..."
        curl -L -o "${IPXE_DIR}/ipxe.efi" "${IPXE_BASE_URL}/ipxe.efi" || {
            log_error "Failed to download ipxe.efi"
            return 1
        }
        log_success "Downloaded ipxe.efi ($(du -h "${IPXE_DIR}/ipxe.efi" | cut -f1))"
    else
        log_info "ipxe.efi already exists, skipping download"
    fi

    # Download UEFI 32-bit bootloader (optional, rare)
    if [[ ! -f "${IPXE_DIR}/ipxe-i386.efi" ]]; then
        log_info "Downloading ipxe-i386.efi (UEFI x86 32-bit bootloader)..."
        curl -L -o "${IPXE_DIR}/ipxe-i386.efi" "${IPXE_BASE_URL}/ipxe-i386.efi" || {
            log_warning "Failed to download ipxe-i386.efi (this is optional)"
        }
        if [[ -f "${IPXE_DIR}/ipxe-i386.efi" ]]; then
            log_success "Downloaded ipxe-i386.efi ($(du -h "${IPXE_DIR}/ipxe-i386.efi" | cut -f1))"
        fi
    else
        log_info "ipxe-i386.efi already exists, skipping download"
    fi

    # Set permissions
    chmod 644 "${IPXE_DIR}"/*.{kpxe,efi} 2>/dev/null || true

    log_success "iPXE bootloaders downloaded successfully"
}

# Build iPXE from source
build_ipxe() {
    log_info "Building iPXE from source..."

    # Check for required tools
    if ! command -v git &> /dev/null; then
        log_error "git is required to build iPXE"
        return 1
    fi

    if ! command -v make &> /dev/null; then
        log_error "make is required to build iPXE"
        return 1
    fi

    # Create temporary build directory
    BUILD_DIR=$(mktemp -d)
    log_info "Build directory: ${BUILD_DIR}"

    # Clone iPXE repository
    log_info "Cloning iPXE repository..."
    git clone https://github.com/ipxe/ipxe.git "${BUILD_DIR}/ipxe" || {
        log_error "Failed to clone iPXE repository"
        return 1
    }

    cd "${BUILD_DIR}/ipxe/src"

    # Build BIOS bootloader
    log_info "Building undionly.kpxe (BIOS bootloader)..."
    make bin/undionly.kpxe || {
        log_error "Failed to build undionly.kpxe"
        return 1
    }
    cp bin/undionly.kpxe "${IPXE_DIR}/undionly.kpxe"
    log_success "Built undionly.kpxe ($(du -h "${IPXE_DIR}/undionly.kpxe" | cut -f1))"

    # Build UEFI bootloader
    log_info "Building ipxe.efi (UEFI x86-64 bootloader)..."
    make bin-x86_64-efi/ipxe.efi || {
        log_error "Failed to build ipxe.efi"
        return 1
    }
    cp bin-x86_64-efi/ipxe.efi "${IPXE_DIR}/ipxe.efi"
    log_success "Built ipxe.efi ($(du -h "${IPXE_DIR}/ipxe.efi" | cut -f1))"

    # Clean up
    cd /
    rm -rf "${BUILD_DIR}"

    # Set permissions
    chmod 644 "${IPXE_DIR}"/*.{kpxe,efi} 2>/dev/null || true

    log_success "iPXE bootloaders built successfully"
}

# Install boot scripts
install_boot_scripts() {
    log_info "Installing boot scripts..."

    # Copy boot.ipxe
    if [[ -f "${SCRIPT_DIR}/ipxe/boot.ipxe" ]]; then
        cp "${SCRIPT_DIR}/ipxe/boot.ipxe" "${IPXE_DIR}/boot.ipxe"
        chmod 644 "${IPXE_DIR}/boot.ipxe"
        log_success "Installed boot.ipxe"
    else
        log_warning "boot.ipxe not found in ${SCRIPT_DIR}/ipxe/"
    fi

    # Copy MAC mappings documentation
    if [[ -f "${SCRIPT_DIR}/ipxe/mac-mappings.txt" ]]; then
        cp "${SCRIPT_DIR}/ipxe/mac-mappings.txt" "${IPXE_DIR}/mac-mappings.txt"
        chmod 644 "${IPXE_DIR}/mac-mappings.txt"
        log_success "Installed mac-mappings.txt"
    fi
}

# Create symlinks for TFTP
create_tftp_symlinks() {
    log_info "Creating TFTP symlinks..."

    # Symlink bootloaders to TFTP directory
    for file in undionly.kpxe ipxe.efi ipxe-i386.efi; do
        if [[ -f "${IPXE_DIR}/${file}" ]]; then
            ln -sf "${IPXE_DIR}/${file}" "${TFTP_DIR}/${file}"
            log_success "Symlinked ${file} to TFTP directory"
        fi
    done
}

# Validate configuration files
validate_configs() {
    log_info "Validating configuration files..."

    local errors=0

    # Check DHCP configuration
    if [[ -f "${SCRIPT_DIR}/dhcp/dhcpd.conf" ]]; then
        log_info "Checking DHCP configuration..."
        if command -v dhcpd &> /dev/null; then
            dhcpd -t -cf "${SCRIPT_DIR}/dhcp/dhcpd.conf" &> /dev/null && \
                log_success "DHCP configuration is valid" || {
                log_error "DHCP configuration is invalid"
                dhcpd -t -cf "${SCRIPT_DIR}/dhcp/dhcpd.conf"
                ((errors++))
            }
        else
            log_warning "dhcpd not installed, skipping DHCP validation"
        fi
    else
        log_error "dhcpd.conf not found"
        ((errors++))
    fi

    # Check Nginx configuration
    if [[ -f "${SCRIPT_DIR}/http/nginx.conf" ]]; then
        log_info "Checking Nginx configuration..."
        if command -v nginx &> /dev/null; then
            nginx -t -c "${SCRIPT_DIR}/http/nginx.conf" &> /dev/null && \
                log_success "Nginx configuration is valid" || {
                log_error "Nginx configuration is invalid"
                nginx -t -c "${SCRIPT_DIR}/http/nginx.conf"
                ((errors++))
            }
        else
            log_warning "nginx not installed, skipping Nginx validation"
        fi
    else
        log_error "nginx.conf not found"
        ((errors++))
    fi

    # Check iPXE boot script
    if [[ -f "${SCRIPT_DIR}/ipxe/boot.ipxe" ]]; then
        log_info "Checking iPXE boot script..."
        # Basic syntax check (iPXE doesn't have a validation tool)
        if grep -q "#!ipxe" "${SCRIPT_DIR}/ipxe/boot.ipxe"; then
            log_success "iPXE boot script appears valid"
        else
            log_error "iPXE boot script is missing #!ipxe shebang"
            ((errors++))
        fi
    else
        log_error "boot.ipxe not found"
        ((errors++))
    fi

    # Check for required bootloaders
    log_info "Checking for iPXE bootloaders..."
    for file in undionly.kpxe ipxe.efi; do
        if [[ -f "${IPXE_DIR}/${file}" ]]; then
            log_success "Found ${file} ($(du -h "${IPXE_DIR}/${file}" | cut -f1))"
        else
            log_warning "${file} not found (run --download or --build-ipxe)"
        fi
    done

    if [[ $errors -eq 0 ]]; then
        log_success "All configuration files are valid"
        return 0
    else
        log_error "Found $errors configuration error(s)"
        return 1
    fi
}

# Test services
test_services() {
    log_info "Testing PXE boot services..."

    local errors=0

    # Test TFTP server
    log_info "Testing TFTP server..."
    if systemctl is-active --quiet atftpd 2>/dev/null; then
        log_success "TFTP server (atftpd) is running"

        # Try to fetch a file via TFTP
        if command -v tftp &> /dev/null; then
            timeout 5 tftp localhost -c get undionly.kpxe /tmp/test-undionly.kpxe &> /dev/null && {
                log_success "TFTP fetch test successful"
                rm -f /tmp/test-undionly.kpxe
            } || {
                log_warning "TFTP fetch test failed (this may be normal if files aren't ready)"
            }
        fi
    else
        log_error "TFTP server is not running"
        ((errors++))
    fi

    # Test HTTP server
    log_info "Testing HTTP server..."
    if systemctl is-active --quiet nginx 2>/dev/null; then
        log_success "HTTP server (nginx) is running"

        # Try to fetch health endpoint
        if command -v curl &> /dev/null; then
            curl -f -s http://localhost/health &> /dev/null && {
                log_success "HTTP health check successful"
            } || {
                log_warning "HTTP health check failed"
                ((errors++))
            }
        fi
    else
        log_error "HTTP server is not running"
        ((errors++))
    fi

    # Test DHCP server
    log_info "Testing DHCP server..."
    if systemctl is-active --quiet dhcpd4 2>/dev/null || \
       systemctl is-active --quiet isc-dhcp-server 2>/dev/null; then
        log_success "DHCP server is running"
    else
        log_error "DHCP server is not running"
        ((errors++))
    fi

    # Network connectivity test
    log_info "Checking network interfaces..."
    ip addr show | grep -q "inet " && {
        log_success "Network interfaces are up"
    } || {
        log_error "No network interfaces with IP addresses found"
        ((errors++))
    }

    if [[ $errors -eq 0 ]]; then
        log_success "All service tests passed"
        return 0
    else
        log_error "Found $errors service error(s)"
        return 1
    fi
}

# Clean up boot assets
clean_assets() {
    log_warning "Cleaning up boot assets directory..."
    read -p "This will delete ${BOOT_ASSETS_DIR}. Continue? (y/N) " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        rm -rf "${BOOT_ASSETS_DIR}"
        rm -rf "${TFTP_DIR}"
        log_success "Boot assets cleaned up"
    else
        log_info "Cleanup cancelled"
    fi
}

# Full installation
full_install() {
    log_info "Starting full PXE server installation..."

    create_directories
    install_boot_scripts
    create_tftp_symlinks

    log_success "Installation complete!"
    log_info ""
    log_info "Next steps:"
    log_info "  1. Download or build iPXE bootloaders:"
    log_info "     sudo $0 --download"
    log_info "     OR"
    log_info "     sudo $0 --build-ipxe"
    log_info ""
    log_info "  2. Configure your network settings in:"
    log_info "     ${SCRIPT_DIR}/dhcp/dhcpd.conf"
    log_info "     ${SCRIPT_DIR}/nixos-module.nix"
    log_info ""
    log_info "  3. Deploy NixOS configuration or manually start services"
    log_info ""
    log_info "  4. Add NixOS boot images to ${NIXOS_DIR}/"
    log_info "     (This will be done by T032.S3 - Image Builder)"
}

# Main script
main() {
    if [[ $# -eq 0 ]]; then
        show_help
        exit 0
    fi

    check_root

    while [[ $# -gt 0 ]]; do
        case $1 in
            --install)
                full_install
                shift
                ;;
            --download)
                download_ipxe
                shift
                ;;
            --build-ipxe)
                build_ipxe
                shift
                ;;
            --validate)
                validate_configs
                shift
                ;;
            --test)
                test_services
                shift
                ;;
            --clean)
                clean_assets
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done
}

main "$@"