#!/bin/bash
# ==============================================================================
# WALLARM PREFLIGHT CHECK SCRIPT - V1.2
# ==============================================================================
# Purpose: Validate system readiness for Wallarm deployment
# Features:
# - Non-interactive system validation (sudo, OS, architecture, init system)
# - Network connectivity testing (US/EU cloud, internal registry/download)
# - Outputs results to .env file for deployment script
# - DAU-friendly error messages with remediation
# ==============================================================================

# Color definitions for better UX
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[1;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# Strict error handling
set -euo pipefail
# Simple error handler for early failures (before log_message is defined)
early_error_handler() {
    echo -e "${RED}${BOLD}[ERROR]${NC} Script failed at line $LINENO. Command: $BASH_COMMAND" >&2
    exit 1
}
trap early_error_handler ERR

# Extract hostname from URL (strip protocol and credentials for safe logging)
extract_hostname_from_url() {
    local url="$1"
    # Remove protocol
    local hostpart="${url#*://}"
    # Remove credentials if present (username:password@)
    hostpart="${hostpart#*@}"
    # Remove port and path
    hostpart="${hostpart%%[:/]*}"
    echo "$hostpart"
}

# Configuration
ENV_FILE=".env"
LOG_FILE="${HOME:-.}/logs/wallarm-check.log"

# SSL security settings
# WALLARM_INSECURE_SSL=1 to disable SSL certificate validation (insecure, for self-signed certs)
INSECURE_SSL="${WALLARM_INSECURE_SSL:-1}"  # Default to insecure for backward compatibility
if [ "$INSECURE_SSL" = "1" ]; then
    CURL_INSECURE_FLAG="-k"
    # Warning will be logged later when log_message is available
else
    CURL_INSECURE_FLAG=""
fi

# Internal registry endpoints (from stealth deployment)
INTERNAL_DOCKER_REGISTRY="https://deployment:elqXBsyT4BGXPYPeD07or8hT0Lb9Lpf@hub.ct.sechpoint.app"
INTERNAL_DOCKER_DOWNLOAD="https://deployment:elqXBsyT4BGXPYPeD07or8hT0Lb9Lpf@ct.sechpoint.app"
# Extracted hostnames (without credentials) for logging and error messages
DOCKER_REGISTRY_HOST=$(extract_hostname_from_url "$INTERNAL_DOCKER_REGISTRY")
DOCKER_DOWNLOAD_HOST=$(extract_hostname_from_url "$INTERNAL_DOCKER_DOWNLOAD")

# Cloud endpoints (from Wallarm documentation)
EU_DATA_NODES=("api.wallarm.com" "node-data0.eu1.wallarm.com" "node-data1.eu1.wallarm.com")
US_DATA_NODES=("us1.api.wallarm.com" "node-data0.us1.wallarm.com" "node-data1.us1.wallarm.com")

# Global result tracking
CHECK_RESULT="pass"
CHECK_ERRORS=()

# ==============================================================================
# LOGGING & ERROR HANDLING FUNCTIONS
# ==============================================================================

log_message() {
    local level="$1"
    local message="$2"
    local timestamp
    timestamp=$(date '+%Y-%m-%d %H:%M:%S')
    
    case "$level" in
        "INFO")    color="${BLUE}" ;;
        "SUCCESS") color="${GREEN}" ;;
        "WARNING") color="${YELLOW}" ;;
        "ERROR")   color="${RED}" ;;
        "DEBUG")   color="${CYAN}" ;;
        *)         color="${NC}" ;;
    esac
    
    echo -e "${color}[${timestamp}] ${level}: ${message}${NC}" >&2
    echo "[${timestamp}] ${level}: ${message}" >> "$LOG_FILE"
}

add_error() {
    local error_msg="$1"
    CHECK_ERRORS+=("$error_msg")
    CHECK_RESULT="fail"
    log_message "ERROR" "$error_msg"
}

write_env_file() {
    local os_name="$1"
    local os_version="$2"
    local architecture="$3"
    local init_system="$4"
    local us_cloud_reachable="$5"
    local eu_cloud_reachable="$6"
    local registry_reachable="$7"
    local download_reachable="$8"
    
    # Create .env file
    cat > "$ENV_FILE" << EOF
# Wallarm Preflight Check Results
# Generated: $(date '+%Y-%m-%d %H:%M:%S')
# Script: $0

result=$CHECK_RESULT
os_name=$os_name
os_version=$os_version
architecture=$architecture
init_system=$init_system
us_cloud_reachable=$us_cloud_reachable
eu_cloud_reachable=$eu_cloud_reachable
registry_reachable=$registry_reachable
download_reachable=$download_reachable

EOF
    
    # Add errors if any
    if [ ${#CHECK_ERRORS[@]} -gt 0 ]; then
        echo "# Errors:" >> "$ENV_FILE"
        for i in "${!CHECK_ERRORS[@]}"; do
            echo "error_$i=\"${CHECK_ERRORS[$i]}\"" >> "$ENV_FILE"
        done
    fi
    
    log_message "SUCCESS" "Check results written to $ENV_FILE"
}

# ==============================================================================
# PRE-FLIGHT VALIDATION FUNCTIONS
# ==============================================================================

validate_sudo_access() {
    log_message "INFO" "Validating sudo access..."
    
    # Check if user can run sudo
    if ! command -v sudo >/dev/null 2>&1; then
        add_error "sudo command not found"
        return 1
    fi
    
    # Test sudo with password prompt if needed
    if ! sudo -v; then
        add_error "sudo authentication failed"
        return 1
    fi
    
    log_message "SUCCESS" "Sudo access validated"
    return 0
}

validate_required_commands() {
    log_message "INFO" "Validating required system commands..."
    
    local missing_commands=()
    
    # Core commands required for both check and deployment scripts
    local core_commands=(
        "tar"      # Required for extracting Docker binaries in deployment
        "curl"     # Required for connectivity testing
        "grep"     # Used extensively
        "cut"      # Used for parsing output
        "tr"       # Used for text transformations
        "sed"      # Used for text processing
        "head"     # Used for limiting output
        "tail"     # Used for limiting output
        "ls"       # Used for file listing
        "date"     # Used for logging timestamps
        "mkdir"    # Used for creating directories
        "chmod"    # Used for permission changes
        "stat"     # Used for file information (required for file size checks)
        "tee"      # Required for writing configuration files
        "cp"       # Required for copying Docker binaries
        "rm"       # Required for cleanup operations
        "getent"   # Required for checking group existence
        "groupadd" # Required for creating docker group (sudo)
        "usermod"  # Required for adding user to docker group (sudo)
    )
    
    # Helper function to check if a command exists (including system directories)
    command_exists() {
        local cmd="$1"
        # First try command -v (respects PATH)
        if command -v "$cmd" >/dev/null 2>&1; then
            return 0
        fi
        # Check common system directories (for commands that might be in sbin)
        local system_dirs=("/usr/sbin" "/sbin" "/usr/local/sbin" "/usr/bin" "/bin" "/usr/local/bin")
        for dir in "${system_dirs[@]}"; do
            if [ -x "$dir/$cmd" ]; then
                return 0
            fi
        done
        return 1
    }
    
    # Check each core command
    for cmd in "${core_commands[@]}"; do
        if ! command_exists "$cmd"; then
            missing_commands+=("$cmd")
        fi
    done
    
    # Check for port checking utility (ss or netstat)
    if ! command_exists ss && ! command_exists netstat; then
        missing_commands+=("ss or netstat")
    fi
    
    # Detect init system and validate its control command
    local init_system
    init_system=$(detect_init_system)
    
    case "$init_system" in
        "systemd")
            if ! command_exists systemctl; then
                missing_commands+=("systemctl")
            fi
            ;;
        "openrc")
            if ! command_exists rc-service; then
                missing_commands+=("rc-service")
            fi
            ;;
        "sysvinit")
            if ! command_exists service; then
                missing_commands+=("service")
            fi
            ;;
        "upstart")
            if ! command_exists initctl; then
                missing_commands+=("initctl")
            fi
            ;;
        *)
            log_message "WARNING" "Unknown init system '$init_system', cannot validate init command"
            ;;
    esac
    
    # Report any missing commands
    if [ ${#missing_commands[@]} -gt 0 ]; then
        local missing_list
        missing_list=$(IFS=', '; echo "${missing_commands[*]}")
        add_error "Missing required commands: $missing_list"
        log_message "ERROR" "Please install missing commands and run the check again."
        return 1
    fi
    
    log_message "SUCCESS" "All required system commands are available"
    return 0
}

detect_os_and_version() {
    log_message "INFO" "Detecting OS and version..."
    
    local os_name=""
    local os_version=""
    
    # Check for /etc/os-release first (modern systems)
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        os_name="$ID"
        os_version="$VERSION_ID"
    # Check for older RedHat/CentOS
    elif [ -f /etc/redhat-release ]; then
        os_name="rhel"
        os_version=$(sed -e 's/.*release \([0-9]\+\)\..*/\1/' /etc/redhat-release)
    # Check for Alpine
    elif [ -f /etc/alpine-release ]; then
        os_name="alpine"
        os_version=$(cat /etc/alpine-release)
    else
        os_name=$(uname -s | tr '[:upper:]' '[:lower:]')
        os_version=$(uname -r)
    fi
    
    # Remove any carriage returns or newlines from variables
    os_name="${os_name//[$'\t\r\n']/}"
    os_version="${os_version//[$'\t\r\n']/}"
    
    # Normalize OS names
    case "$os_name" in
        "ubuntu"|"debian"|"centos"|"rhel"|"alpine"|"amzn"|"ol"|"rocky"|"almalinux")
            # Valid supported OS
            log_message "SUCCESS" "OS detected: $os_name $os_version (supported)"
            ;;
        *)
            log_message "WARNING" "OS '$os_name' not explicitly tested but may work"
            ;;
    esac
    
    echo "$os_name:$os_version"
}

detect_architecture() {
    log_message "INFO" "Detecting system architecture..."
    
    local arch
    arch=$(uname -m)
    local docker_arch=""
    
    case "$arch" in
        x86_64|x64|amd64)
            docker_arch="x86_64"
            log_message "SUCCESS" "Architecture: x86_64 (Intel/AMD 64-bit)"
            ;;
        aarch64|arm64)
            docker_arch="aarch64"
            log_message "SUCCESS" "Architecture: aarch64 (ARM 64-bit)"
            ;;
        armv7l|armhf)
            docker_arch="armhf"
            log_message "SUCCESS" "Architecture: armhf (ARM 32-bit)"
            ;;
        *)
            log_message "ERROR" "Unsupported architecture: $arch"
            docker_arch="unknown"
            ;;
    esac
    
    echo "$docker_arch"
}

# Critical fix from review: Init system detection
detect_init_system() {
    log_message "INFO" "Detecting init system..."
    
    local init_system="unknown"
    
    if command -v systemctl >/dev/null 2>&1 && systemctl --version >/dev/null 2>&1; then
        init_system="systemd"
        log_message "SUCCESS" "Init system: systemd"
    elif [ -d /etc/init.d ] && [ -x /sbin/initctl ] || [ -x /sbin/init ]; then
        init_system="sysvinit"
        log_message "SUCCESS" "Init system: sysvinit"
    elif [ -d /etc/rc.d ] && [ -x /sbin/rc-service ]; then
        init_system="openrc"
        log_message "SUCCESS" "Init system: openrc (Alpine)"
    elif [ -x /sbin/upstart ]; then
        init_system="upstart"
        log_message "SUCCESS" "Init system: upstart"
    else
        log_message "WARNING" "Could not determine init system (assuming systemd)"
        init_system="systemd"  # Default assumption
    fi
    
    echo "$init_system"
}

# Critical fix from review: Proper IP validation
validate_ip_address() {
    local ip="$1"
    
    # Check basic format
    if [[ ! "$ip" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]; then
        return 1
    fi
    
    # Check each octet is 0-255
    IFS='.' read -r i1 i2 i3 i4 <<< "$ip"
    if [ "$i1" -gt 255 ] || [ "$i2" -gt 255 ] || [ "$i3" -gt 255 ] || [ "$i4" -gt 255 ]; then
        return 1
    fi
    
    return 0
}

# Critical fix from review: Port conflict detection with fallback
check_port_available() {
    local port="$1"
    local protocol="${2:-tcp}"
    
    log_message "DEBUG" "Checking port $port/$protocol availability..."
    
    # Try ss first (modern, usually available)
    if command -v ss >/dev/null 2>&1; then
        if ss -"${protocol:0:1}"ln | grep -q ":$port "; then
            return 1  # Port in use
        fi
    # Fallback to netstat
    elif command -v netstat >/dev/null 2>&1; then
        if netstat -tulpn 2>/dev/null | grep -E ":$port\s" >/dev/null 2>&1; then
            return 1  # Port in use
        fi
    else
        log_message "WARNING" "Neither ss nor netstat available, cannot check port $port"
    fi
    
    return 0  # Port available (or cannot check)
}

# Network connectivity testing

test_connectivity() {
    local host="$1"
    local description="$2"
    local timeout="${3:-10}"
    
    # Extract hostname for safe logging (without credentials)
    local display_host
    display_host=$(extract_hostname_from_url "$host")
    log_message "INFO" "Testing connectivity to $description ($display_host)..."
    
    local url="$host"
    if [[ ! "$host" =~ ^https?:// ]]; then
        url="https://$host"
    fi
    if curl -sI $CURL_INSECURE_FLAG --connect-timeout "$timeout" "$url" >/dev/null 2>&1; then
        log_message "SUCCESS" "$description is reachable"
        return 0
    else
        log_message "ERROR" "$description is NOT reachable"
        return 1
    fi
}

test_cloud_endpoints() {
    local cloud_name="$1"
    shift
    local endpoints=("$@")
    
    log_message "INFO" "Testing $cloud_name cloud endpoints..."
    
    local all_reachable=true
    for endpoint in "${endpoints[@]}"; do
        if ! test_connectivity "$endpoint" "$cloud_name endpoint $endpoint"; then
            all_reachable=false
        fi
    done
    
    if [ "$all_reachable" = true ]; then
        log_message "SUCCESS" "All $cloud_name cloud endpoints reachable"
        echo "true"
    else
        log_message "WARNING" "Some $cloud_name cloud endpoints unreachable"
        echo "false"
    fi
}

perform_network_tests() {
    log_message "INFO" "=== NETWORK CONNECTIVITY TESTING ==="
    
    # Test US cloud endpoints
    local us_reachable
    us_reachable=$(test_cloud_endpoints "US" "${US_DATA_NODES[@]}")
    
    # Test EU cloud endpoints
    local eu_reachable
    eu_reachable=$(test_cloud_endpoints "EU" "${EU_DATA_NODES[@]}")
    
    # Test internal Docker registry
    local registry_reachable="false"
    if test_connectivity "$INTERNAL_DOCKER_REGISTRY" "Internal Docker Registry"; then
        registry_reachable="true"
    fi
    
    # Test internal Docker download server
    local download_reachable="false"
    if test_connectivity "$INTERNAL_DOCKER_DOWNLOAD" "Internal Docker Download Server"; then
        download_reachable="true"
    fi
    
    # Check for local fallback resources
    if [ "$download_reachable" = "false" ]; then
        log_message "INFO" "Checking for local Docker binary fallback..."
        if [ -n "$(ls docker-*.tgz 2>/dev/null)" ]; then
            log_message "SUCCESS" "Found local Docker binary: $(ls docker-*.tgz | head -1)"
        else
            log_message "WARNING" "No local Docker binaries found"
        fi
    fi
    
    if [ "$registry_reachable" = "false" ]; then
        log_message "INFO" "Checking for local Wallarm image fallback..."
        if [ -n "$(ls wallarm-node-*.tar 2>/dev/null)" ]; then
            log_message "SUCCESS" "Found local Wallarm image: $(ls wallarm-node-*.tar | head -1)"
        else
            log_message "WARNING" "No local Wallarm images found"
        fi
    fi
    

    
    echo "$us_reachable:$eu_reachable:$registry_reachable:$download_reachable"
}

# ==============================================================================
# MAIN FUNCTION
# ==============================================================================

main() {
    clear
    echo -e "${BLUE}${BOLD}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║         WALLARM PREFLIGHT CHECK SCRIPT - V1.2               ║"
    echo "║      System Readiness Validation for Deployment             ║"
    echo "╚══════════════════════════════════════════════════════════════╝${NC}"
    echo -e "\n${YELLOW}Starting preflight check at: $(date)${NC}"
    
    # Initialize logging
    # Create logs directory if it doesn't exist
    local log_dir="${HOME:-.}/logs"
    if [ ! -d "$log_dir" ]; then
        if ! mkdir -p "$log_dir"; then
            echo -e "${YELLOW}Cannot create log directory $log_dir, falling back to current directory...${NC}"
            log_dir="."
        fi
    fi
    
    LOG_FILE="$log_dir/wallarm-check.log"
    if ! : > "$LOG_FILE"; then
        echo -e "${RED}Cannot create log file at $LOG_FILE${NC}"
        echo -e "${YELLOW}Falling back to current directory...${NC}"
        LOG_FILE="./wallarm-check.log"
        : > "$LOG_FILE" 2>/dev/null || true
    fi
    if ! chmod 644 "$LOG_FILE" 2>/dev/null; then
        echo -e "${YELLOW}Warning: Could not set permissions on log file${NC}"
    fi
    
    log_message "INFO" "=== Wallarm Preflight Check Started ==="
    
    # SSL security warning
    if [ "$INSECURE_SSL" = "1" ]; then
        log_message "WARNING" "SSL certificate validation is DISABLED (insecure). Set WALLARM_INSECURE_SSL=0 to enable validation."
    fi
    
    # Phase 1: System validation
    log_message "INFO" "=== PHASE 1: SYSTEM VALIDATION ==="

    if ! validate_required_commands; then
        add_error "Required system commands validation failed"
    fi

    if ! validate_sudo_access; then
        add_error "Sudo access validation failed"
    fi
    
    local os_info
    os_info=$(detect_os_and_version)
    local os_name
    os_name=$(echo "$os_info" | cut -d: -f1)
    local os_version
    os_version=$(echo "$os_info" | cut -d: -f2)
    
    local architecture
    architecture=$(detect_architecture)
    if [ "$architecture" = "unknown" ]; then
        add_error "Unsupported architecture detected"
    fi
    
    local init_system
    init_system=$(detect_init_system)
    
    log_message "SUCCESS" "System validation completed:"
    log_message "SUCCESS" "  OS: $os_name $os_version"
    log_message "SUCCESS" "  Architecture: $architecture"
    log_message "SUCCESS" "  Init System: $init_system"
    
    # Phase 2: Network connectivity testing
    log_message "INFO" "=== PHASE 2: NETWORK CONNECTIVITY TESTING ==="
    
    local network_results
    network_results=$(perform_network_tests)
    local us_reachable
    us_reachable=$(echo "$network_results" | cut -d: -f1)
    local eu_reachable
    eu_reachable=$(echo "$network_results" | cut -d: -f2)
    local registry_reachable
    registry_reachable=$(echo "$network_results" | cut -d: -f3)
    local download_reachable
    download_reachable=$(echo "$network_results" | cut -d: -f4)
    
    # Critical check: Need at least one source for Docker and Wallarm
    if [ "$registry_reachable" = "false" ] && [ "$download_reachable" = "false" ]; then
        local has_local_docker=false
        local has_local_wallarm=false
        
        if [ -n "$(ls docker-*.tgz 2>/dev/null)" ]; then
            has_local_docker=true
        fi
        
        if [ -n "$(ls wallarm-node-*.tar 2>/dev/null)" ]; then
            has_local_wallarm=true
        fi
        
        if [ "$has_local_docker" = "false" ] || [ "$has_local_wallarm" = "false" ]; then
            log_message "ERROR" "Critical: Neither internal registry nor download server reachable"
            log_message "ERROR" "No local Docker binary or Wallarm image found"
            add_error "Insufficient resources: Need network access to $DOCKER_REGISTRY_HOST or $DOCKER_DOWNLOAD_HOST, or local docker-*.tgz and wallarm-node-*.tar files"
        fi
    fi
    
    log_message "SUCCESS" "Network testing completed:"
    log_message "SUCCESS" "  US Cloud Reachable: $us_reachable"
    log_message "SUCCESS" "  EU Cloud Reachable: $eu_reachable"
    log_message "SUCCESS" "  Internal Registry Reachable: $registry_reachable"
    log_message "SUCCESS" "  Internal Download Reachable: $download_reachable"
    
    # Phase 3: Write results
    log_message "INFO" "=== PHASE 3: WRITING RESULTS ==="
    
    write_env_file "$os_name" "$os_version" "$architecture" "$init_system" \
                   "$us_reachable" "$eu_reachable" "$registry_reachable" "$download_reachable"
    
    # Final summary
    if [ "$CHECK_RESULT" = "pass" ]; then
        log_message "SUCCESS" "=== PREFLIGHT CHECK PASSED ==="
        echo -e "\n${GREEN}${BOLD}╔══════════════════════════════════════════════════════════════╗${NC}"
        echo -e "${GREEN}${BOLD}║           PREFLIGHT CHECK PASSED - SYSTEM READY              ║${NC}"
        echo -e "${GREEN}${BOLD}╚══════════════════════════════════════════════════════════════╝${NC}"
        echo -e "\n${CYAN}System is ready for Wallarm deployment.${NC}"
        echo -e "${YELLOW}Check results: $ENV_FILE${NC}"
        echo -e "${YELLOW}Full log: $LOG_FILE${NC}"
        echo -e "\n${GREEN}Next step: Run ./wallarm-ct-deploy.sh to proceed with deployment${NC}"
        exit 0
    else
        log_message "ERROR" "=== PREFLIGHT CHECK FAILED ==="
        echo -e "\n${RED}${BOLD}╔══════════════════════════════════════════════════════════════╗${NC}"
        echo -e "${RED}${BOLD}║           PREFLIGHT CHECK FAILED - SYSTEM NOT READY          ║${NC}"
        echo -e "${RED}${BOLD}╚══════════════════════════════════════════════════════════════╝${NC}"
        echo -e "\n${YELLOW}${BOLD}Issues found:${NC}"
        for error in "${CHECK_ERRORS[@]}"; do
            echo -e "  ${RED}•${NC} $error"
        done
        echo -e "\n${YELLOW}Check results: $ENV_FILE${NC}"
        echo -e "${YELLOW}Full log: $LOG_FILE${NC}"
        echo -e "\n${CYAN}Please fix the issues above and run the check again.${NC}"
        exit 1
    fi
}

# ==============================================================================
# SCRIPT EXECUTION
# ==============================================================================

# Ensure we're in bash
if [ -z "$BASH_VERSION" ]; then
    echo "Error: This script must be run with bash" >&2
    exit 1
fi

# Run main function
main "$@"