#!/usr/bin/env bash
set -euo pipefail

#####################################
# BOOTSTRAP - Download required files
#####################################
SCRIPT_SOURCE="${BASH_SOURCE[0]}"
SCRIPT_DIR="$(cd "$(dirname "$SCRIPT_SOURCE")" && pwd)"

# Server URL - update this when server changes
REBOOTLESS_SERVER_URL="${REBOOTLESS_SERVER_URL:-https://ampps.nuftp.com/rebootless_src}"

bootstrap_rebootless() {
    local required_files=("rebootless.conf")
    local required_dirs=("scripts" "dashboard")
    local missing=0
    
    echo "[BOOTSTRAP] Checking for required files in: $SCRIPT_DIR"
    
    for file in "${required_files[@]}"; do
        if [[ ! -f "$SCRIPT_DIR/$file" ]]; then
            echo "[BOOTSTRAP] Missing: $file"
            missing=1
        fi
    done
    
    for dir in "${required_dirs[@]}"; do
        if [[ ! -d "$SCRIPT_DIR/$dir" ]]; then
            echo "[BOOTSTRAP] Missing: $dir/"
            missing=1
        fi
    done
    
    if [[ $missing -eq 0 && "${REBOOTLESS_UPDATE:-0}" != "1" ]]; then
        echo "[BOOTSTRAP] All required files present"
        return 0
    fi
    
    if [[ "${REBOOTLESS_UPDATE:-0}" == "1" ]]; then
        echo "[BOOTSTRAP] Force update enabled - re-downloading files..."
    else
        echo "[BOOTSTRAP] Downloading required files from server..."
    fi
    echo "[BOOTSTRAP] Server: $REBOOTLESS_SERVER_URL"
    
    cd "$SCRIPT_DIR"
    
    # echo "[BOOTSTRAP] Downloading rebootless.conf..."
    # curl -fsSL -o rebootless.conf "$REBOOTLESS_SERVER_URL/rebootless.conf" || {
    #     echo "[BOOTSTRAP] ERROR: Failed to download rebootless.conf"
    #     exit 1
    # }
    
    echo "[BOOTSTRAP] Downloading scripts/..."
    rm -rf scripts rebootless.tgz
    curl -fsSL -o rebootless.tgz "$REBOOTLESS_SERVER_URL/rebootless.tgz" || {
        echo "[BOOTSTRAP] ERROR: Failed to download scripts"
        exit 1
    }
    tar -xzf rebootless.tgz && rm rebootless.tgz
    
    echo "[BOOTSTRAP] Download complete!"
}

bootstrap_rebootless

#####################################
# CONFIG
#####################################
KERNEL_VERSION="$(uname -r | sed 's/\.x86_64$//')"

# Auto-detect Linux distribution
detect_distro() {
    if [[ -f /etc/almalinux-release ]]; then
        echo "almalinux"
    elif [[ -f /etc/rhel-release ]] || [[ -f /etc/redhat-release ]]; then
        echo "rhel"
    elif [[ -f /etc/centos-release ]]; then
        echo "centos"
    elif [[ -f /etc/rocky-release ]]; then
        echo "rocky"
    elif [[ -f /etc/debian_version ]]; then
        echo "debian"
    elif [[ -f /etc/ubuntu-release ]]; then
        echo "ubuntu"
    else
        echo "unknown"
    fi
}

DISTRO=$(detect_distro)
echo "Detected distribution: $DISTRO"

# Auto-detect livepatch support
check_livepatch_support() {
    if [[ -f /proc/config.gz ]]; then
        zcat /proc/config.gz | grep -q "CONFIG_LIVEPATCH=y" && echo "1" || echo "0"
    elif [[ -f /boot/config-$(uname -r) ]]; then
        grep -q "CONFIG_LIVEPATCH=y" /boot/config-$(uname -r) && echo "1" || echo "0"
    elif [[ -s /sys/kernel/livepatch ]]; then
        echo "1"
    else
        echo "0"
    fi
}

LIVEPATCH_SUPPORT=$(check_livepatch_support)
if [[ "$LIVEPATCH_SUPPORT" == "0" ]]; then
    echo -e "${YELLOW}WARNING: Livepatch not enabled in kernel${RESET}"
fi

# Load config file if exists
CONFIG_FILE="${CONFIG_FILE:-$SCRIPT_DIR/rebootless.conf}"
if [[ -f "$CONFIG_FILE" ]]; then
    source "$CONFIG_FILE"
fi

# Allow env vars to override config
BASE_DIR="${BASE_DIR:-$SCRIPT_DIR}"
DATA_DIR="${DATA_DIR:-$BASE_DIR/rebootlessdata}"
PATCH_DIR="${PATCH_DIR:-$BASE_DIR/kernel_cve_patches}"
SRC_RPM="${SRC_RPM:-$BASE_DIR/kernel-${KERNEL_VERSION}.src.rpm}"
KERNEL_SRC_DIR="${KERNEL_SRC_DIR:-$BASE_DIR/source}"
SCRIPTS_DIR="${SCRIPTS_DIR:-$BASE_DIR/scripts}"
LOGFILE="${LOGFILE:-$BASE_DIR/rebootless_automation.log}"
SERVER_BASE_URL="${SERVER_BASE_URL:-https://ampps.nuftp.com/rebootless_jsons/}"
MAX_PARALLEL_DOWNLOADS="${MAX_PARALLEL_DOWNLOADS:-5}"
MAX_PARALLEL_BUILDS="${MAX_PARALLEL_BUILDS:-2}"
SKIP_EXISTING_MODULES="${SKIP_EXISTING_MODULES:-1}"
TEST_CUMULATIVE_BEFORE_BUILD="${TEST_CUMULATIVE_BEFORE_BUILD:-1}"

# Auto-enable prereq install if running as root
if [[ -z "${AUTO_INSTALL_PREREQS:-}" ]]; then
    if [[ $EUID -eq 0 ]]; then
        AUTO_INSTALL_PREREQS=1
    else
        AUTO_INSTALL_PREREQS=0
    fi
fi

# Resume tracking
STATE_FILE="$DATA_DIR/.automation_state"

#####################################
# LOGGING
#####################################
exec > >(tee -a "$LOGFILE") 2>&1

YELLOW="\e[33m"
CYAN="\e[36m"
GREEN="\e[32m"
RED="\e[31m"
MAGENTA="\e[35m"
RESET="\e[0m"
NC="${RESET}"
#####################################
# HELP
#####################################
usage() {
cat <<EOF
Usage:
  $0 <step> [step...]

Bootstrap:
  On first run, scripts are automatically downloaded from server.
  Set REBOOTLESS_SERVER_URL to use a custom server.

Steps:
  prereqs                           Install prerequisites kpatch and required libs
  get_cves                          Run cves_from_updateinfo.py
  get_patches                       Download CVE patches
  check                             Detect empty CVE patch directories
  kernel-src                        Extract kernel SRPM and source
  score                             Score and select best patches
  vmlinux                           Build vmlinux
  kpatch                            Build kpatch modules (Test patches)
  build_module <MODULE_VERSION>     Build livepatch module
  pre                               all steps except Build kpatch
  all                               Run everything
  resume                            Resume from last completed step
  status                            Show current status

Options:
  REBOOTLESS_SERVER_URL=<url>       Custom server URL for downloading files
  RESUME=0                          Disable resume (run all steps)
  SKIP_EXISTING_MODULES=0           Rebuild all modules

Examples:
  $0 all v1
  $0 cves patches score
  $0 vmlinux
  $0 build_module v5
  RESUME=0 $0 all v1                # Run without resume
  REBOOTLESS_SERVER_URL=http://myserver $0 all v1  # Custom server
EOF
exit 1
}

#####################################
# STEPS
#####################################
install_prerequisite(){
    echo "[STEP] Checking prerequisites"
    
    #####################################
    # Must be root
    #####################################
    if [[ $EUID -ne 0 ]]; then
        echo "ERROR: Must be run as root"
        exit 1
    fi
    
    #####################################
    # Detect package manager
    #####################################
    local PKG_MGR=""
    local PKG_CHECK_CMD=""
    local PKG_INSTALL_CMD=""
    
    if command -v dnf >/dev/null 2>&1; then
        PKG_MGR="dnf"
        PKG_CHECK_CMD="rpm -q"
        PKG_INSTALL_CMD="dnf install -y"
    elif command -v yum >/dev/null 2>&1; then
        PKG_MGR="yum"
        PKG_CHECK_CMD="rpm -q"
        PKG_INSTALL_CMD="yum install -y"
    elif command -v apt-get >/dev/null 2>&1; then
        PKG_MGR="apt"
        PKG_CHECK_CMD="dpkg -s"
        PKG_INSTALL_CMD="apt-get install -y"
    else
        echo "ERROR: No supported package manager found (dnf/yum/apt)"
        exit 1
    fi
    
    echo "Detected package manager: $PKG_MGR"
    
    #####################################
    # Define required packages by distro
    #####################################
    local MISSING_PKGS=()
    local REQUIRED_PKGS=()
    
    if [[ "$PKG_MGR" == "dnf" || "$PKG_MGR" == "yum" ]]; then
        REQUIRED_PKGS=(
            gcc
            make
            binutils
            elfutils-libelf-devel
            rpm-build
            cpio
            tar
            xz
            git
            openssl
            bison
            flex
            openssl-devel
            bc
            dwarves
            php
            php-json
            jq
            curl
        )
    else  # apt-based
        REQUIRED_PKGS=(
            gcc
            make
            binutils
            libelf-dev
            cpio
            tar
            xz-utils
            git
            openssl
            bison
            flex
            libssl-dev
            bc
            dwarves
            php
            jq
            curl
        )
    fi
    
    #####################################
    # Collect missing packages
    #####################################
    for pkg in "${REQUIRED_PKGS[@]}"; do
        if ! $PKG_CHECK_CMD "$pkg" >/dev/null 2>&1; then
            MISSING_PKGS+=("$pkg")
        fi
    done
    
    #####################################
    # Kernel livepatch support (non-package)
    #####################################
    local KERNEL_CONFIG=""
    if [[ -f "/boot/config-$(uname -r)" ]]; then
        KERNEL_CONFIG="/boot/config-$(uname -r)"
    elif [[ -f "/proc/config.gz" ]]; then
        KERNEL_CONFIG="/proc/config.gz"
    fi
    
    if [[ -n "$KERNEL_CONFIG" ]]; then
        if [[ "$KERNEL_CONFIG" == *.gz ]]; then
            if ! zcat "$KERNEL_CONFIG" | grep -q "CONFIG_LIVEPATCH=y"; then
                echo "ERROR: Running kernel does not have CONFIG_LIVEPATCH enabled"
                exit 1
            fi
        else
            if ! grep -q "CONFIG_LIVEPATCH=y" "$KERNEL_CONFIG"; then
                echo "ERROR: Running kernel does not have CONFIG_LIVEPATCH enabled"
                exit 1
            fi
        fi
    else
        echo "WARNING: Could not find kernel config to verify CONFIG_LIVEPATCH"
    fi
    
    #####################################
    # Show and install missing packages
    #####################################
    if (( ${#MISSING_PKGS[@]} > 0 )); then
        echo
        echo "The following prerequisite packages are missing:"
        echo "----------------------------------------------"
        for pkg in "${MISSING_PKGS[@]}"; do
            echo " - $pkg"
        done
        
        confirm_or_exit "These packages are required to continue."
        
        echo
        echo "Installing missing packages..."
        
        # Update package cache for apt
        if [[ "$PKG_MGR" == "apt" ]]; then
            apt-get update
        fi
        
        $PKG_INSTALL_CMD "${MISSING_PKGS[@]}"
    else
        echo "OK: All required packages are installed"
    fi
    
    #####################################
    # Disk space warning (non-fatal)
    #####################################
    local avail
    avail=$(df --output=avail "$BASE_DIR" | tail -1)
    if (( avail < 10000000 )); then
        echo "WARNING: Less than 10GB free disk space"
    fi
    
    #####################################
    # kpatch check/build
    #####################################
    if command -v kpatch >/dev/null 2>&1; then
        echo "OK: kpatch already installed"
        kpatch --version || true
        return
    fi
    
    echo
    confirm_or_exit "kpatch is not installed. It will be built from source."
    
    local KPATCH_DIR="$BASE_DIR/kpatch"
    if [[ ! -d "$KPATCH_DIR/.git" ]]; then
        git clone https://github.com/dynup/kpatch.git "$KPATCH_DIR"
    fi
    
    cd "$KPATCH_DIR"
    make
    make install
    
    #####################################
    # Final verification
    #####################################
    if ! command -v kpatch >/dev/null 2>&1; then
        echo "ERROR: kpatch installation failed"
        exit 1
    fi
    
    echo "kpatch installed successfully"
    kpatch --version || true
    set_state "prereqs"
}

extract_cves_from_text() {
    grep -oE 'CVE-[0-9]{4}-[0-9]+' "$DATA_DIR/cves_$KERNEL_VERSION.txt" | sort -u
}

extract_cves_from_json() {
    local json="$1"
    jq -r 'keys[]' "$json" | sort
}

compare_cves_text_vs_json() {
    local kernel="$KERNEL_VERSION"
    local text_file="$DATA_DIR/cves_$KERNEL_VERSION.txt"
    local json_file="$DATA_DIR/${kernel}_cve_index_ser.json"
    local server_url="$SERVER_BASE_URL/${kernel}_cve_index.json"

    # Ensure text file exists
    if [[ ! -f "$text_file" ]]; then
        echo -e "${RED}❌ CVE text file not found${RESET}"
        echo -e "${CYAN}Please run : $0 get_cves${RESET}"
        return 1
    fi

    # If JSON missing locally, try downloading
    if curl --head --silent --fail "$server_url" >/dev/null; then
        echo -e "${GREEN}Downloading latest JSON...${RESET}"

        if ! curl -s "$server_url" -o "$json_file"; then
            echo -e "${RED}❌ ERROR: JSON download failed.${RESET}"
            exit 1
        fi
    else
        echo "Checking for local JSON copy"

        if [[ ! -f "$json_file" ]]; then
            echo -e "${YELLOW}⚠️ No local JSON available. Assuming all CVEs are new.${RESET}"
            return 0
        else
            echo -e "${YELLOW}ℹ️ Using local JSON.${RESET}"
        fi
    fi

    # Extract CVEs
    text_cves=$(extract_cves_from_text)
    json_cves=$(extract_cves_from_json "$json_file")

    # Compare
    new_cves=$(comm -13 <(echo "$json_cves") <(echo "$text_cves"))

    if [[ -n "$new_cves" ]]; then
        echo -e "${GREEN}New CVEs detected:${MAGENTA}"
        echo "$new_cves"
        echo -e "${RESET}"
        return 0
    else
        return 1
    fi
}

spinner() {
    local pid=$1
    local msg="$2"
    local delay=0.1
    local spin=('⠋' '⠙' '⠹' '⠸' '⠼' '⠴' '⠦' '⠧' '⠇' '⠏')

    exec 3>/dev/tty || return   # Open terminal FD

    tput civis >&3 2>/dev/null

    while kill -0 "$pid" 2>/dev/null; do
        for c in "${spin[@]}"; do
            printf "\r${CYAN}%s${RESET} %s" "$c" "$msg" >&3
            sleep $delay
        done
    done

    printf "\r" >&3
    tput cnorm >&3 2>/dev/null

    exec 3>&-   # Close FD
}

#####################################
# STATE MANAGEMENT (Resume capability)
#####################################
get_state() {
    local step="$1"
    if [[ -f "$STATE_FILE" ]]; then
        grep "^${step}=" "$STATE_FILE" 2>/dev/null | cut -d= -f2-
    fi
}

set_state() {
    local step="$1"
    local value="${2:-done}"
    mkdir -p "$(dirname "$STATE_FILE")"
    sed -i "/^${step}=/d" "$STATE_FILE" 2>/dev/null || true
    echo "${step}=${value}" >> "$STATE_FILE"
}

clear_state() {
    rm -f "$STATE_FILE"
}

is_step_completed() {
    local step="$1"
    [[ "$(get_state "$step")" == "done" ]]
}

should_skip_step() {
    local step="$1"
    local skip_file="$DATA_DIR/.skip_${step}"
    
    if [[ -f "$skip_file" ]]; then
        echo -e "${YELLOW}Skipping $step (disabled)${RESET}"
        return 0
    fi
    
    if [[ "${RESUME:-1}" -eq 1 ]] && is_step_completed "$step"; then
        echo -e "${GREEN}Skipping $step (already completed)${RESET}"
        return 0
    fi
    
    return 1
}

#####################################
# PROGRESS TRACKING
#####################################
total_steps=0
current_step=0

init_progress() {
    total_steps=$1
    current_step=0
    echo -e "${CYAN}========================================${RESET}"
    echo -e "${CYAN}Starting automation: $total_steps steps${RESET}"
    echo -e "${CYAN}========================================${RESET}"
}

update_progress() {
    local step_name="$1"
    ((current_step++))
    echo -e "${CYAN}[${current_step}/${total_steps}] ${step_name}${RESET}"
}

#####################################
# DOWNLOAD PROGRESS
#####################################
download_progress() {
    local current=$1
    local total=$2
    local percent=$((current * 100 / total))
    printf "\r${CYAN}Progress: [%d/%d] %d%%${RESET}" "$current" "$total" "$percent"
}

confirm_or_exit() {
    local message="$1"

    if [[ "${AUTO_INSTALL_PREREQS:-0}" -eq 1 ]]; then
        echo "AUTO_INSTALL_PREREQS=1 set — proceeding automatically"
        return 0
    fi

    echo
    echo "$message"
    read -r -p "Do you want to continue? (yes/no): " answer

    case "$answer" in
        [Yy]|[Yy][Ee][Ss]) return 0 ;;
        *) echo "Aborting."; exit 1 ;;
    esac
}

find_cves_list() {
    clear
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║               [STEP] Extracting CVEs                     ║"
    echo "╚══════════════════════════════════════════════════════════╝"

    # Detect OS family and call appropriate function
    detect_os
    
    case "$OS_FAMILY" in
        debian)
            find_cves_list_ubuntu
            ;;
        rhel)
            find_cves_list_rhel
            ;;
    esac
    
    [[ -s "$DATA_DIR/cves_$KERNEL_VERSION.txt" ]] \
        || { echo -e "${RED}ERROR: CVE list empty${RESET}"; exit 1; }

    if compare_cves_text_vs_json; then
        # New CVEs found (return 0) - continue automatically
        echo -e "${GREEN}New CVEs found. Continuing...${RESET}"
    else
        # No new CVEs (return 1) - ask user
        echo -e "${GREEN}✅ No new CVEs found.${RESET}"
        echo -e "${YELLOW}Do you want to continue anyway? (y/N, timeout 10s)${RESET}"

        if read -t 10 -r -n 1 answer; then
            echo
            case "$answer" in
                y|Y)
                    echo -e "${GREEN}Continuing...${RESET}"
                    ;;
                *)
                    echo -e "${RED}Exiting.${RESET}"
                    exit 0
                    ;;
            esac
        else
            echo
            echo -e "${RED}No input received. Exiting.${RESET}"
            exit 0
        fi
    fi
    set_state "get_cves"
}

#####################################
# Get CVEs - Ubuntu/Debian (Like yum updateinfo)
#####################################
find_cves_list_ubuntu() {
    echo -e "${CYAN}Detected Ubuntu/Debian - using Ubuntu CVE Finder (yum updateinfo equiv)${RESET}"
    
    # For Ubuntu, use running kernel version directly
    KERNEL_VERSION="$(uname -r)"
    
    # Try new ubuntu_updateinfo.sh first, fallback to ubuntu_kernel_cve_finder.sh
    local ubuntu_updateinfo="$SCRIPTS_DIR/ubuntu_updateinfo.sh"
    local ubuntu_finder="$SCRIPTS_DIR/ubuntu_kernel_cve_finder.sh"
    
    if [[ -f "$ubuntu_updateinfo" ]]; then
        echo -e "${CYAN}Using: ubuntu_updateinfo.sh${RESET}"
        "$ubuntu_updateinfo" "$KERNEL_VERSION"
    elif [[ -f "$ubuntu_finder" ]]; then
        echo -e "${CYAN}Using: ubuntu_kernel_cve_finder.sh${RESET}"
        "$ubuntu_finder"
    else
        echo -e "${RED}ERROR: No Ubuntu CVE finder script found in $SCRIPTS_DIR${RESET}"
        exit 1
    fi
    
    # Verify CVE list was created
    local cve_list_file="$DATA_DIR/cves_needed_for_livepatch.txt"
    if [[ -f "$cve_list_file" && -s "$cve_list_file" ]]; then
        local cve_count=$(grep -c 'CVE' "$cve_list_file" || echo 0)
        echo -e "${GREEN}✅ Found $cve_count CVEs needing patches${RESET}"
        echo -e "${GREEN}✅ CVE list saved to: $cve_list_file${RESET}"
    else
        echo -e "${RED}ERROR: Failed to generate CVE list${RESET}"
        exit 1
    fi
}

#####################################
# Get CVEs - RHEL/CentOS/AlmaLinux
#####################################
find_cves_list_rhel() {
    echo -e "${CYAN}Detected RHEL/CentOS/AlmaLinux - using yum updateinfo${RESET}"
    
    php "$SCRIPTS_DIR/cves_from_updateinfo.php" 

    [[ -s "$DATA_DIR/cves_$KERNEL_VERSION.txt" ]] \
        || { echo -e "${RED}ERROR: CVE list empty${RESET}"; exit 1; }
}

download_patches() {
    detect_os
    
    if [[ "$OS_FAMILY" == "debian" ]]; then
        if [[ ! -f "$DATA_DIR/cves_$KERNEL_VERSION.txt" || ! -s "$DATA_DIR/cves_$KERNEL_VERSION.txt" ]]; then
            echo -e "${RED}❌ CVE text file not found or empty${RESET}"
            echo -e "${CYAN}Please run : $0 get_cves${RESET}"
            return 1
        fi
        
        clear
        echo "╔══════════════════════════════════════════════════════════╗"
        echo "║             [STEP] Downloading Patches                   ║"
        echo "╚══════════════════════════════════════════════════════════╝"

        php "$SCRIPTS_DIR/kernel-cve-kpatch.php" \
            "$DATA_DIR/cves_$KERNEL_VERSION.txt"

        STATUS=$?

        if [[ $STATUS -ne 0 ]]; then
            echo -e "❌ Patch download failed. Check log: $LOGFILE"
            exit 1
        fi
        
        local patch_count=$(find "$DATA_DIR/kernel_cve_patches" -name "*.patch" 2>/dev/null | wc -l)
        local cve_count=$(ls -d "$DATA_DIR/kernel_cve_patches"/CVE-* 2>/dev/null | wc -l)
        
        echo ""
        echo -e "Found ${GREEN}$patch_count${NC} patches for ${GREEN}$cve_count${NC} CVEs"
        
        # Update CVE index with available patches
        echo -e "${CYAN}Updating CVE index...${NC}"
        php -r '
        $data_dir = getenv("DATA_DIR") ?: "rebootlessdata";
        $kernel_version = getenv("KERNEL_VERSION") ?: trim(shell_exec("uname -r"));
        $patch_dir = $data_dir . "/kernel_cve_patches";
        
        $data = [];
        if (is_dir($patch_dir)) {
            $dirs = glob($patch_dir . "/CVE-*");
            foreach ($dirs as $dir) {
                $cve = basename($dir);
                $patches = glob($dir . "/*.patch");
                $valid_patches = [];
                foreach ($patches as $patch) {
                    $size = filesize($patch);
                    if ($size > 100) {
                        $handle = fopen($patch, "r");
                        $header = fread($handle, 10);
                        fclose($handle);
                        if (strpos($header, "<!DOCTYPE") === false) {
                            $valid_patches[] = basename($patch);
                        }
                    }
                }
                if (!empty($valid_patches)) {
                    $data[$cve] = [
                        "cve" => $cve,
                        "cvss_v2" => null,
                        "cvss_v3" => null,
                        "description" => "",
                        "cve_url" => "https://nvd.nist.gov/vuln/detail/" . $cve,
                        "selected_patch" => null,
                        "available_patches" => []
                    ];
                }
            }
        }
        
        $json_file = $data_dir . "/" . $kernel_version . "_cve_index.json";
        file_put_contents($json_file, json_encode($data, JSON_PRETTY_PRINT));
        echo "Created CVE index with " . count($data) . " CVEs\n";
        '
        
        echo -e "✅ KernelCare patches ready for scoring."
        set_state "get_patches"
    
    elif [[ "$OS_FAMILY" == "rhel" ]]; then
        if [[ ! -f "$DATA_DIR/cves_$KERNEL_VERSION.txt" || ! -s "$DATA_DIR/cves_$KERNEL_VERSION.txt" ]]; then
            echo -e "${RED}❌ CVE text file not found or empty${RESET}"
            echo -e "${CYAN}Please run : $0 get_cves${RESET}"
            return 1
        fi
        
        clear
        echo "╔══════════════════════════════════════════════════════════╗"
        echo "║             [STEP] Downloading Patches                   ║"
        echo "╚══════════════════════════════════════════════════════════╝"

        php "$SCRIPTS_DIR/kernel-cve-kpatch.php" \
            "$DATA_DIR/cves_$KERNEL_VERSION.txt"

        STATUS=$?

        if [[ $STATUS -ne 0 ]]; then
            echo -e "❌ Patch download failed. Check log: $LOGFILE"
            exit 1
        fi
        
        local patch_count=$(find "$DATA_DIR/kernel_cve_patches" -name "*.patch" 2>/dev/null | wc -l)
        local cve_count=$(ls -d "$DATA_DIR/kernel_cve_patches"/CVE-* 2>/dev/null | wc -l)
        
        echo ""
        echo -e "Found ${GREEN}$patch_count${NC} patches for ${GREEN}$cve_count${NC} CVEs"
        
        echo -e "${CYAN}Updating CVE index...${NC}"
        php -r '
        $data_dir = getenv("DATA_DIR") ?: "rebootlessdata";
        $kernel_version = getenv("KERNEL_VERSION") ?: trim(shell_exec("uname -r"));
        $patch_dir = $data_dir . "/kernel_cve_patches";
        
        $data = [];
        if (is_dir($patch_dir)) {
            $dirs = glob($patch_dir . "/CVE-*");
            foreach ($dirs as $dir) {
                $cve = basename($dir);
                $patches = glob($dir . "/*.patch");
                $valid_patches = [];
                foreach ($patches as $patch) {
                    $size = filesize($patch);
                    if ($size > 100) {
                        $handle = fopen($patch, "r");
                        $header = fread($handle, 10);
                        fclose($handle);
                        if (strpos($header, "<!DOCTYPE") === false) {
                            $valid_patches[] = basename($patch);
                        }
                    }
                }
                if (!empty($valid_patches)) {
                    $data[$cve] = [
                        "cve" => $cve,
                        "cvss_v2" => null,
                        "cvss_v3" => null,
                        "description" => "",
                        "cve_url" => "https://nvd.nist.gov/vuln/detail/" . $cve,
                        "selected_patch" => null,
                        "available_patches" => []
                    ];
                }
            }
        }
        
        $json_file = $data_dir . "/" . $kernel_version . "_cve_index.json";
        file_put_contents($json_file, json_encode($data, JSON_PRETTY_PRINT));
        echo "Created CVE index with " . count($data) . " CVEs\n";
        '
        
        echo -e "✅ KernelCare patches ready for scoring."
        set_state "get_patches"
    fi
}


check_empty_folders() {
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║          [STEP] Checking empty CVE directories           ║"
    echo "╚══════════════════════════════════════════════════════════╝"

    find "$PATCH_DIR" -type d -empty || true
    set_state "check"
}

detect_os() {
    if [[ -f /etc/os-release ]]; then
        # Save colors before source — os-release can overwrite variables
        local _yellow="$YELLOW" _cyan="$CYAN" _green="$GREEN"
        local _red="$RED" _magenta="$MAGENTA" _reset="$RESET"

        source /etc/os-release

        # Restore colors
        YELLOW="$_yellow"; CYAN="$_cyan"; GREEN="$_green"
        RED="$_red"; MAGENTA="$_magenta"; RESET="$_reset"

        OS_ID="${ID,,}"
        OS_VERSION="$VERSION_ID"
    else
        echo -e "${RED}ERROR: Cannot detect OS. /etc/os-release not found.${RESET}"
        exit 1
    fi

    case "$OS_ID" in
        ubuntu|debian)                          OS_FAMILY="debian" ;;
        almalinux|rhel|centos|rocky|fedora)     OS_FAMILY="rhel"   ;;
        *)
            echo -e "${RED}ERROR: Unsupported OS: $OS_ID${RESET}"
            exit 1
            ;;
    esac

    echo -e "${CYAN}Detected OS     : $OS_ID $OS_VERSION (family: $OS_FAMILY)${RESET}"
}

#####################################
# Debian: resolve full kernel version
# uname gives 5.4.0-29-generic
# we need  5.4.0-29.33 for Launchpad
#####################################
detect_debian_kernel_version() {
    local running_kernel
    running_kernel=$(uname -r)

    local pkg_version

    # Step 1: dpkg — cleanest and local
    pkg_version=$(dpkg -l "linux-image-${running_kernel}" 2>/dev/null \
        | awk '/^ii/ {print $3}' | head -1)
    [[ -n "$pkg_version" ]] && echo -e "${CYAN}Version source  : dpkg${RESET}"

    # Step 2: /proc/version_signature — Ubuntu specific
    if [[ -z "$pkg_version" ]] && [[ -f /proc/version_signature ]]; then
        pkg_version=$(grep -oP '\d+\.\d+\.\d+-\d+\.\d+' /proc/version_signature | head -1)
        [[ -n "$pkg_version" ]] && echo -e "${CYAN}Version source  : /proc/version_signature${RESET}"
    fi

    # Step 3: apt-cache policy
    if [[ -z "$pkg_version" ]]; then
        pkg_version=$(apt-cache policy "linux-image-${running_kernel}" 2>/dev/null \
            | awk '/Installed:/ {print $2}' \
            | grep -oP '^\d+\.\d+\.\d+-\d+\.\d+')
        [[ -n "$pkg_version" ]] && echo -e "${CYAN}Version source  : apt-cache policy${RESET}"
    fi

    # Step 4: Last resort — stripped uname (no patch number)
    if [[ -z "$pkg_version" ]]; then
        pkg_version=$(echo "$running_kernel" \
            | sed 's/-generic//;s/-lowlatency//;s/-aws//;s/-azure//;s/-gcp//;s/-kvm//')
        echo -e "${YELLOW}WARN: Could not resolve full package version, using $pkg_version${RESET}"
        echo -e "${YELLOW}HINT: Manually set KERNEL_VERSION=x.x.x-xx.xx if download fails${RESET}"
    fi

    # Override KERNEL_VERSION with full Debian package version
    KERNEL_VERSION="$pkg_version"
    KERNEL_BASE=$(echo "$pkg_version" | grep -oP '^\d+\.\d+\.\d+')

    echo -e "${GREEN}Kernel version  : $KERNEL_VERSION${RESET}"
    echo -e "${GREEN}Kernel base     : $KERNEL_BASE${RESET}"
}

#####################################
# Get Kernel Source — RHEL/AlmaLinux
#####################################
get_kernel_src_rhel() {
    local tarball="linux-${KERNEL_VERSION}.tar.xz"

    # 1. Check if kernel tarball exists
    if [[ -f "$BASE_DIR/$tarball" ]]; then
        echo -e "${GREEN}OK: Kernel tarball already exists: $tarball${RESET}"
    else
        echo -e "${YELLOW}Kernel tarball not found, extracting from SRPM...${RESET}"

        if [[ ! -f "$SRC_RPM" ]]; then
            echo -e "${YELLOW}SRPM not found at $SRC_RPM, downloading from upstream vault...${RESET}"

            # --- Derive OS version from kernel version string ---
            # e.g. 4.18.0-477.10.1.el8_8 → 8.8
            local os_version
            os_version=$(echo "$KERNEL_VERSION" | grep -oP 'el\d+_\d+' | grep -oP '\d+_\d+' | tr '_' '.')

            if [[ -z "$os_version" ]]; then
                echo -e "${RED}ERROR: Could not derive OS version from KERNEL_VERSION='$KERNEL_VERSION'${RESET}"
                echo -e "${YELLOW}Expected format e.g.: 4.18.0-477.10.1.el8_8${RESET}"
                exit 1
            fi

            local major_version="${os_version%%.*}"

            # --- Detect distro and build vault URL ---
            local vault_base
            local distro_id distro_version_id

            if [[ -f /etc/os-release ]]; then
                distro_id=$(grep -oP '^ID=\K[^"]+' /etc/os-release | tr -d '"' | tr '[:upper:]' '[:lower:]')
                distro_version_id=$(grep -oP '^VERSION_ID=\K[^"]+' /etc/os-release | tr -d '"')
            fi

            case "$distro_id" in
                almalinux|alma)
                    vault_base="https://vault.almalinux.org/${os_version}/BaseOS/Source/Packages"
                    ;;
                centos)
                    # CentOS 8+ uses vault.centos.org with a different layout
                    if [[ "$major_version" -ge 8 ]]; then
                        vault_base="https://vault.centos.org/${os_version}/BaseOS/Source/Packages"
                    else
                        # CentOS 7 and older: flat layout
                        vault_base="https://vault.centos.org/${os_version}/os/Source/SPackages"
                    fi
                    ;;
                rhel|"red hat enterprise linux")
                    # No public SRPM vault; fall back to dnf download
                    echo -e "${YELLOW}RHEL detected — falling back to 'dnf download --source'...${RESET}"
                    dnf download --source kernel-"$KERNEL_VERSION" \
                        --destdir "$BASE_DIR" >> "$LOGFILE" 2>&1
                    SRC_RPM=$(ls "$BASE_DIR"/kernel-*.src.rpm 2>/dev/null | head -1)
                    ;;
                rocky|rockylinux)
                    vault_base="https://download.rockylinux.org/vault/rocky/${os_version}/BaseOS/source/tree/Packages/k"
                    ;;
                *)
                    echo -e "${RED}ERROR: Unsupported or unrecognised distro '${distro_id}'. Cannot determine SRPM vault URL.${RESET}"
                    exit 1
                    ;;
            esac

            # --- Download SRPM from vault (skip if dnf fallback already ran) ---
            if [[ -z "$SRC_RPM" || ! -f "$SRC_RPM" ]]; then
                echo -e "${CYAN}Searching for kernel SRPM at: $vault_base${RESET}"

                local srpm_filename
                srpm_filename=$(curl -fsSL "$vault_base/" \
                    | grep -oP 'kernel-[^"]+\.src\.rpm' \
                    | grep "$KERNEL_VERSION" \
                    | head -1)

                if [[ -z "$srpm_filename" ]]; then
                    echo -e "${RED}ERROR: Could not find kernel SRPM for '$KERNEL_VERSION' at $vault_base${RESET}"
                    exit 1
                fi

                local srpm_url="${vault_base}/${srpm_filename}"
                echo -e "${CYAN}Downloading SRPM: $srpm_url${RESET}"

                curl -fL --progress-bar "$srpm_url" -o "$BASE_DIR/$srpm_filename" \
                    >> "$LOGFILE" 2>&1

                if [[ $? -ne 0 ]]; then
                    echo -e "${RED}ERROR: Failed to download SRPM from $srpm_url${RESET}"
                    exit 1
                fi

                SRC_RPM="$BASE_DIR/$srpm_filename"
            fi
        fi

        if [[ -z "$SRC_RPM" || ! -f "$SRC_RPM" ]]; then
            echo -e "${RED}ERROR: Could not find or download kernel SRPM${RESET}"
            exit 1
        fi

        mkdir -p "$BASE_DIR/kernel-source"
        cd "$BASE_DIR/kernel-source"
        rpm2cpio "$SRC_RPM" | cpio -idmv >> "$LOGFILE" 2>&1

        if [[ ! -f "$tarball" ]]; then
            echo -e "${RED}ERROR: $tarball not found inside SRPM${RESET}"
            exit 1
        fi

        cp "$tarball" "$BASE_DIR/"
    fi

    # 2. Extract kernel source if needed
    if [[ -d "$KERNEL_SRC_DIR" ]] && [[ -n "$(ls -A "$KERNEL_SRC_DIR")" ]]; then
        echo -e "${GREEN}Kernel source directory already exists: $KERNEL_SRC_DIR${RESET}"
        return
    fi

    echo -e "${CYAN}Extracting kernel source...${RESET}"
    cd "$BASE_DIR"
    tar -xf "$tarball" >> "$LOGFILE" 2>&1

    if [[ ! -d "linux-${KERNEL_VERSION}" ]]; then
        echo -e "${RED}ERROR: Extracted kernel directory missing${RESET}"
        exit 1
    fi

    cp -a "linux-${KERNEL_VERSION}" "$KERNEL_SRC_DIR"
    echo -e "${GREEN}Kernel source prepared successfully${RESET}"
}

#####################################
# Get Kernel Source — Ubuntu/Debian
#####################################
get_kernel_src_debian() {
    local orig_tarball="linux_${KERNEL_BASE}.orig.tar.gz"
    local diff_gz="linux_${KERNEL_VERSION}.diff.gz"
    local dsc_file="linux_${KERNEL_VERSION}.dsc"
    local base_url="https://launchpad.net/ubuntu/+archive/primary/+sourcefiles/linux/${KERNEL_VERSION}"

    # 1. Check/Download source files
    cd "$BASE_DIR"
    for file in "$orig_tarball" "$diff_gz" "$dsc_file"; do
        if [[ -f "$BASE_DIR/$file" ]]; then
            echo -e "${GREEN}OK: Already exists: $file${RESET}"
        else
            echo -e "${CYAN}Downloading: $file ...${RESET}"
            wget -q --show-progress -P "$BASE_DIR" "${base_url}/${file}"
            if [[ $? -ne 0 ]]; then
                echo -e "${RED}ERROR: Failed to download $file${RESET}"
                echo -e "${YELLOW}HINT: Check https://launchpad.net/ubuntu/+source/linux for valid versions${RESET}"
                exit 1
            fi
        fi
    done

    # 2. Extract kernel source if needed
    if [[ -d "$KERNEL_SRC_DIR" ]] && [[ -n "$(ls -A "$KERNEL_SRC_DIR")" ]]; then
        echo -e "${GREEN}Kernel source directory already exists: $KERNEL_SRC_DIR${RESET}"
        return
    fi

    echo -e "${CYAN}Extracting kernel source with dpkg-source...${RESET}"
    cd "$BASE_DIR"
    dpkg-source -x "$dsc_file" "$KERNEL_SRC_DIR" >> "$LOGFILE" 2>&1

    if [[ ! -d "$KERNEL_SRC_DIR" ]]; then
        echo -e "${RED}ERROR: Extracted kernel directory missing${RESET}"
        exit 1
    fi

    echo -e "${GREEN}Kernel source prepared successfully${RESET}"
}

#####################################
# Main
#####################################
get_kernel_src() {
    clear
    echo -e "${MAGENTA}╔══════════════════════════════════════════════════════════╗${RESET}"
    echo -e "${MAGENTA}║       [STEP] Preparing Kernel Source                     ║${RESET}"
    echo -e "${MAGENTA}╚══════════════════════════════════════════════════════════╝${RESET}"

    detect_os

    case "$OS_FAMILY" in
        rhel)
            echo -e "${GREEN}Kernel version  : $KERNEL_VERSION${RESET}"
            get_kernel_src_rhel
            ;;
        debian)
            detect_debian_kernel_version
            get_kernel_src_debian
            ;;
    esac
    set_state "kernel-src"
}

find_best_patches() {
    clear
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║             echo "[STEP] Scoring patches"                ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    
    #python3 "$SCRIPTS_DIR/score_and_select_patches.py" \
    #    "$KERNEL_SRC_DIR" "$KERNEL_VERSION"

    php "$SCRIPTS_DIR/score_select_patches.php" "$KERNEL_SRC_DIR" "$KERNEL_VERSION"

    local PATCH_OUT="${KERNEL_VERSION}_patches"

    [[ -d "$PATCH_OUT" ]] || {
        echo "ERROR: Patch output directory not found: $PATCH_OUT"
        exit 1
    }

    echo "╔═════════════════════════════════════════════════════════════════════╗"
    echo "║     [STEP] Checking for CVE directories with multiple patches       ║"
    echo "╚═════════════════════════════════════════════════════════════════════╝"

    local MULTI_PATCH_DIRS=()
    while IFS= read -r dir; do
        MULTI_PATCH_DIRS+=("$dir")
    done < <(
        find "$PATCH_OUT" -mindepth 1 -maxdepth 1 -type d \
        -exec bash -c '(( $(ls -1 "$1" | wc -l) > 1 ))' _ {} \; \
        -print
    )

    if (( ${#MULTI_PATCH_DIRS[@]} > 0 )); then
        echo
        echo "WARNING: Multiple patches detected for the following CVEs:"
        echo "----------------------------------------------------------"
        for dir in "${MULTI_PATCH_DIRS[@]}"; do
            echo "CVE: $(basename "$dir")"
            ls -1 "$dir"
            echo
        done

        #####################################
        # Confirmation gate
        #####################################
        if [[ "${AUTO_APPROVE_MULTI_PATCHES:-0}" -eq 1 ]]; then
            echo "AUTO_APPROVE_MULTI_PATCHES=1 set — proceeding automatically"
        else
            echo "Do you want to continue with these CVEs? (yes/no)"
            read -r answer

            case "$answer" in
                [Nn]|[Nn][Oo])
                    echo "Aborting due to multiple patches per CVE"
                    exit 1
                ;;
                [Yy]|[Yy][Ee][Ss])
                    echo "Proceeding as requested"
                ;;
        *)
            echo "Invalid response. Please answer yes or no."
            exit 1
            ;;
    esac
fi
    else
        echo "OK: No CVE directories with multiple patches"
    fi
    set_state "score"
}

make_vmlinux() {
    clear
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║              [STEP] Building vmlinux                     ║"
    echo "╚══════════════════════════════════════════════════════════╝"

    cd "$KERNEL_SRC_DIR" || {
        echo -e "${RED}Failed to cd into $KERNEL_SRC_DIR${RESET}"
        exit 1
    }
    if [ ! -f vmlinux ];then
        cp "/boot/config-$(uname -r)" .config

        openssl req -new -x509 -nodes -days 3650 \
            -subj "/CN=kernel-build-rhel/" \
            -out certs/rhel.pem >> "$BASE_DIR/rebootless_automation.log" 2>&1 || exit 1

        make olddefconfig >> "$BASE_DIR/rebootless_automation.log" 2>&1 || exit 1
        echo -e "${GREEN}Kernel configuration prepared.${RESET}"
        
        make vmlinux -j"$(nproc)" >> "$BASE_DIR/rebootless_automation.log" 2>&1 &

        PID=$!
        spinner $PID "Building vmlinux... "

        wait $PID
        STATUS=$?

        if [[ $STATUS -ne 0 || ! -f vmlinux ]]; then
            echo -e "${RED}\r❌ vmlinux build failed.... Check log: $LOGFILE ${RESET}"
            exit 1
        else
            echo -e "${GREEN}\r✅  vmlinux built successfully.${RESET}"
        fi
    else
        echo -e "${GREEN}vmlinux already present in source.${RESET}"
    fi
    set_state "vmlinux"
}


check_patches_to_modules() {
    clear
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║           [STEP] Building kpatch modules                 ║"
    echo "╚══════════════════════════════════════════════════════════╝"

    "$SCRIPTS_DIR/v2_build_kpatch_modules.sh"
    set_state "kpatch"
}

check_version_arg(){

    if [[ $# -lt 2 ]]; then
        echo -e " "
        echo -e "${CYAN}  Please run $0 build_module <module_version> for final livepatch module."
        echo -e "${CYAN}  $0 build_module <module_version> "
        echo -e "${RESET} "
        exit 0;
    fi
}

build_module(){
    clear
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║           [STEP] Building livepatch module               ║"
    echo "╚══════════════════════════════════════════════════════════╝"

    if [[ $# -eq 2 ]]; then
        local kernel_version="$1"
        local module_version="$2"

        if [[ -z "$kernel_version" ]]; then
            echo "Usage: $0 build_module <kernel_version> <module_version>"
            exit 1
        fi

        php $SCRIPTS_DIR/collect_successful_patches.php

        if [[ -n "$KERNEL_SRC_DIR" && "$KERNEL_SRC_DIR" != "/" ]]; then
            rm -rf "$KERNEL_SRC_DIR"
        else
            echo "Refusing to delete: $KERNEL_SRC_DIR"
        fi

        "$SCRIPTS_DIR/patches_to_livepatch.sh" "$kernel_version" "$module_version"
    else
        echo "Usage: $0 build_module <kernel_version> <module_version>"
        exit 1
    fi
}

#####################################
# DISPATCH
#####################################
[[ $# -eq 0 ]] && usage

STEP="$1"

case "$STEP" in
    prereqs)     install_prerequisite ;;
    get_cves)    find_cves_list ;;
    get_patches) download_patches ;;
    check)       check_empty_folders ;;
    kernel-src)  get_kernel_src ;;
    score)       detect_os
                  find_best_patches ;;
    vmlinux)     detect_os
                  make_vmlinux ;;
    kpatch)      detect_os
                  check_patches_to_modules ;;

    build_module)
        if [[ $# -eq 2 ]]; then
            MODULE_VERSION="$2"

            if [[ -z "$KERNEL_VERSION" ]]; then
                echo "Usage: $0 build_module <kernel_version> <module_version>"
                exit 1
            fi

            build_module "$KERNEL_VERSION" "$MODULE_VERSION"
        else
            echo "Usage: $0 build_module <kernel_version> <module_version>"
            exit 1
        fi
        ;;

    pre)
        install_prerequisite
        find_cves_list
        download_patches
        check_empty_folders
        get_kernel_src
        find_best_patches
        make_vmlinux
        ;;

    all)
        if [[ $# -lt 2 ]]; then
            echo "Usage: $0 all <module_version>"
            echo "Example: $0 all v1"
            exit 1
        fi
        MODULE_VERSION="$2"
        
        install_prerequisite
        find_cves_list
        download_patches
        check_empty_folders
        get_kernel_src
        find_best_patches
        make_vmlinux
        check_patches_to_modules
        build_module "$KERNEL_VERSION" "$MODULE_VERSION"
        ;;

    status)
        show_status
        ;;

    resume)
        run_resume
        ;;

    *) usage ;;
esac

#####################################
# STATUS & RESUME FUNCTIONS
#####################################
show_status() {
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║               Automation Status                          ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    
    if [[ ! -f "$STATE_FILE" ]]; then
        echo -e "${YELLOW}No automation has been run yet.${RESET}"
        return
    fi
    
    echo -e "${CYAN}Completed steps:${RESET}"
    while IFS='=' read -r step status; do
        echo -e "  ✓ ${step}: ${status}"
    done < "$STATE_FILE"
    
    echo ""
    echo -e "${CYAN}To resume from last step:${RESET}"
    echo -e "  $0 resume"
    echo ""
    echo -e "${CYAN}To start fresh:${RESET}"
    echo -e "  RESUME=0 $0 all v1"
}

run_resume() {
    if [[ ! -f "$STATE_FILE" ]]; then
        echo -e "${RED}No resume point found. Run a full workflow first.${RESET}"
        exit 1
    fi
    
    echo -e "${CYAN}Resuming from last completed step...${RESET}"
    
    # Read last completed step and continue
    last_step=$(tail -1 "$STATE_FILE" | cut -d= -f1)
    
    case "$last_step" in
        prereqs)
            find_cves_list
            ;;
        get_cves)
            download_patches
            ;;
        get_patches)
            check_empty_folders
            get_kernel_src
            find_best_patches
            make_vmlinux
            check_patches_to_modules
            ;;
        kernel-src)
            find_best_patches
            make_vmlinux
            check_patches_to_modules
            ;;
        score)
            make_vmlinux
            check_patches_to_modules
            ;;
        vmlinux)
            check_patches_to_modules
            ;;
        *)
            echo -e "${YELLOW}Resume not supported from '$last_step'. Starting fresh.${RESET}"
            clear_state
            ;;
    esac
}


echo "══════════════════════════════ Completed: $* ══════════════════════════════"