#!/usr/bin/env bash
#####################################
# Rebootless Rollback Manager
#####################################

set -euo pipefail

BASE_DIR="/root/rebootless"
ROLLBACK_DIR="$BASE_DIR/rollback"

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
RESET='\033[0m'

usage() {
    cat <<EOF
Usage: $0 [COMMAND] [OPTIONS]

Commands:
  list                    List available rollback modules
  load <module>          Load a specific rollback module
  latest                 Load the latest rollback module
  info <module>          Show module information
  clean                  Remove old rollback modules (keeps last 3)

Examples:
  $0 list
  $0 latest
  $0 load RL-4.18.0-513.9.1.el8_9-v1-20260101.ko
  $0 clean
EOF
    exit 0
}

list_modules() {
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║            Available Rollback Modules                    ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    echo ""
    
    if [[ ! -d "$ROLLBACK_DIR" ]] || [[ -z "$(ls -A "$ROLLBACK_DIR" 2>/dev/null)" ]]; then
        echo -e "${YELLOW}No rollback modules found${RESET}"
        return
    fi
    
    echo -e "${CYAN}Rollback directory: $ROLLBACK_DIR${RESET}"
    echo ""
    
    local idx=1
    for mod in $(ls -t "$ROLLBACK_DIR"/RL-*.ko 2>/dev/null); do
        local name=$(basename "$mod")
        local size=$(ls -lh "$mod" | awk '{print $5}')
        local date=$(stat -c %y "$mod" 2>/dev/null | cut -d' ' -f1 || echo "unknown")
        
        echo -e "${GREEN}$idx)${RESET} $name"
        echo -e "    Size: $size | Date: $date"
        echo ""
        ((idx++)) || true
    done
    
    echo -e "Total: $((idx-1)) module(s)"
    echo ""
    echo -e "To load a module:"
    echo -e "  $0 load <module_name>"
    echo -e "  $0 latest"
}

load_module() {
    local module="$1"
    
    if [[ -z "$module" ]]; then
        echo -e "${RED}Error: Module name required${RESET}"
        exit 1
    fi
    
    local mod_path="$ROLLBACK_DIR/$module"
    
    if [[ ! -f "$mod_path" ]]; then
        echo -e "${RED}Error: Module not found: $module${RESET}"
        echo "Use '$0 list' to see available modules"
        exit 1
    fi
    
    echo -e "${CYAN}Loading module: $module${RESET}"
    
    # Check if running as root
    if [[ $EUID -ne 0 ]]; then
        echo -e "${YELLOW}Warning: Not running as root. Use sudo to load module.${RESET}"
    fi
    
    # Unload current livepatch if any
    if command -v kpatch >/dev/null 2>&1; then
        local current=$(kpatch list 2>/dev/null | grep -v "^Loaded" | head -1 || true)
        if [[ -n "$current" ]]; then
            echo -e "${YELLOW}Unloading current livepatch: $current${RESET}"
            sudo kpatch unload "$current" 2>/dev/null || true
        fi
    fi
    
    # Load the module
    echo -e "${CYAN}Loading module...${RESET}"
    sudo insmod "$mod_path"
    
    if [[ $? -eq 0 ]]; then
        echo -e "${GREEN}✓ Module loaded successfully${RESET}"
        
        # Verify
        if command -v kpatch >/dev/null 2>&1; then
            echo ""
            echo -e "${CYAN}Current livepatches:${RESET}"
            kpatch list 2>/dev/null || echo "  (kpatch list not available)"
        fi
    else
        echo -e "${RED}✗ Failed to load module${RESET}"
        exit 1
    fi
}

load_latest() {
    local latest=$(ls -t "$ROLLBACK_DIR"/RL-*.ko 2>/dev/null | head -1)
    
    if [[ -z "$latest" ]]; then
        echo -e "${RED}No rollback modules found${RESET}"
        exit 1
    fi
    
    local name=$(basename "$latest")
    echo -e "${CYAN}Loading latest module: $name${RESET}"
    load_module "$name"
}

show_info() {
    local module="$1"
    
    if [[ -z "$module" ]]; then
        echo -e "${RED}Error: Module name required${RESET}"
        exit 1
    fi
    
    local mod_path="$ROLLBACK_DIR/$module"
    
    if [[ ! -f "$mod_path" ]]; then
        echo -e "${RED}Error: Module not found: $module${RESET}"
        exit 1
    fi
    
    echo "╔══════════════════════════════════════════════════════════╗"
    echo "║                 Module Information                         ║"
    echo "╚══════════════════════════════════════════════════════════╝"
    echo ""
    
    modinfo "$mod_path" 2>/dev/null || echo "modinfo not available"
}

clean_old() {
    echo -e "${CYAN}Cleaning old rollback modules (keeping last 3)...${RESET}"
    
    local total=$(ls -1 "$ROLLBACK_DIR"/RL-*.ko 2>/dev/null | wc -l)
    
    if [[ $total -le 3 ]]; then
        echo -e "${YELLOW}Already have 3 or fewer modules, nothing to clean${RESET}"
        return
    fi
    
    local removed=0
    for old_mod in $(ls -t "$ROLLBACK_DIR"/RL-*.ko 2>/dev/null | tail -n +4); do
        rm -f "$old_mod" 2>/dev/null || true
        ((removed++)) || true
    done
    
    echo -e "${GREEN}Removed $removed old module(s)${RESET}"
    
    echo ""
    list_modules
}

# Main
COMMAND="${1:-}"

case "$COMMAND" in
    list)
        list_modules
        ;;
    load)
        load_module "$2"
        ;;
    latest)
        load_latest
        ;;
    info)
        show_info "$2"
        ;;
    clean)
        clean_old
        ;;
    -h|--help|help)
        usage
        ;;
    *)
        usage
        ;;
esac