#!/bin/bash

. colors

VERS="0.6.0"

HUGES_RX=1168
HUGES_CN=128
MEM_RES_PERC=90                                                                 # % of available memory for reservation

HPAGES_REQ=$1                                                                   # Requested number of HugePages

MEM_AVAIL=`cat /proc/meminfo | grep "MemFree" | awk '{ print $2 }'`             # Memory available
HPAGE_SIZE=`cat /proc/meminfo | grep "Hugepagesize" | awk '{ print $2 }'`       # HugePage size usually 2M
HPAGES_CUR=`cat /proc/meminfo | grep "HugePages_Total" | awk '{ print $2 }'`    # Number of total 2Mb HugePages
HPAGES_FREE=`cat /proc/meminfo | grep "HugePages_Free" | awk '{ print $2 }'`    # Number of free 2Mb HugePages
HPAGES_MAX=$(($HPAGES_FREE+$MEM_AVAIL/$HPAGE_SIZE*$MEM_RES_PERC/100))           # Number of HugePages (max calculated)

CPU_NUMA=`lscpu | grep "NUMA node(s)" | awk '{ print $3 }'`
CPU_CORE=`lscpu | grep "^CPU(s):" | awk '{ print $2 }'`

HPAGES_1GB_MEM=$(( $CPU_NUMA * 3 * 1024 * 1024 + 2048 * $CPU_CORE ))

HUGES_RX_CALC=$(( $HUGES_RX * $CPU_NUMA + $CPU_CORE ))

huge_rx_1gb_pages_get(){
    HPAGES_1GB_FREE=0
    HPAGES_1GB_TOTAL=0
    for i in $(find /sys/devices/system/node/node* -maxdepth 0 -type d);
    do
        [[ -d "$i/hugepages/hugepages-1048576kB" ]] && HPAGES_1GB_FREE=$(($HPAGES_1GB_FREE+`cat $i/hugepages/hugepages-1048576kB/free_hugepages`)) && HPAGES_1GB_TOTAL=$(($HPAGES_1GB_TOTAL+`cat $i/hugepages/hugepages-1048576kB/nr_hugepages`))
    done
}

huge_rx_1gb_pages_set(){
   local qty=
   [[ -z $1 ]] && qty=3 || qty=$1
   for i in $(find /sys/devices/system/node/node* -maxdepth 0 -type d);
   do
       echo $qty > "$i/hugepages/hugepages-1048576kB/nr_hugepages";
   done
}

memory_clear1(){
   echo -e "${GREEN}[INFO]${NOCOLOR} Reset PageCache"
   sync; echo 1 > /proc/sys/vm/drop_caches
}

memory_clear2(){
   echo -e "${GREEN}[INFO]${NOCOLOR} Reset denries and inodes"
   sync; echo 2 > /proc/sys/vm/drop_caches
}

memory_clear3(){
   echo -e "${GREEN}[INFO]${NOCOLOR} Reset PageCache, denries and inodes"
   sync; echo 3 > /proc/sys/vm/drop_caches
}

randomx_intel_boost(){
    cpu_check=`lscpu | grep "Model name" | grep -c "Intel"`
    msr_check=`lscpu | grep "^Flags" | grep -c "msr"`
    if [[ "$cpu_check" -gt 0 && "$msr_check" -gt 0 ]]; then
        pkg_check=`dpkg -l | grep -c "msr-tools"`
        if [[ "$pkg_check" -eq 0 ]]; then
            echo "Installing msr-tools package ..."
            apt install -y msr-tools || return
        fi
        [[ `lsmod | grep msr` == "" ]] && modprobe msr
        wrmsr -a 0x1a4 0xf
        echo -e "${GREEN}[INFO]${NOCOLOR} Intel CPU RandomX boost applied"
    fi
}

randomx_ryzen_boost(){
    cpu_check=`lscpu | grep "Model name" | grep -c -E "AMD Ryzen|AMD EPYC"`
    msr_check=`lscpu | grep "^Flags" | grep -c "msr"`
    if [[ "$cpu_check" -gt 0 && "$msr_check" -gt 0 ]]; then
        pkg_check=`dpkg -l | grep -c "msr-tools"`
        if [[ "$pkg_check" -eq 0 ]]; then
            echo "Installing msr-tools package ..."
            apt install -y msr-tools || return
        fi
        [[ `lsmod | grep msr` == "" ]] && modprobe msr || return
        
        # Get CPU family and model for architecture detection
        CPU_FAMILY=`cat /proc/cpuinfo | grep "cpu family" | head -n1 | awk '{print $4}'`
        CPU_MODEL=`cat /proc/cpuinfo | grep "model[[:space:]]:" | head -n1 | awk '{print $3}'`
        
        if [[ "$CPU_FAMILY" == "26" ]]; then
            echo -e "${GREEN}[INFO]${NOCOLOR} Detected Ryzen (Zen5)"
            # Zen5 MSR values - using optimized values for Granite Ridge/Strix Point
            wrmsr -a 0xc0011020 0x4400000000000 2>/dev/null || echo -e "${YELLOW}[WARN]${NOCOLOR} MSR 0xc0011020 write failed"
            wrmsr -a 0xc0011021 0x4000000000040 2>/dev/null || echo -e "${YELLOW}[WARN]${NOCOLOR} MSR 0xc0011021 write failed"
            wrmsr -a 0xc0011022 0x8680000401570000 2>/dev/null || echo -e "${YELLOW}[WARN]${NOCOLOR} MSR 0xc0011022 write failed"
            wrmsr -a 0xc001102b 0x2040cc10 2>/dev/null || echo -e "${YELLOW}[WARN]${NOCOLOR} MSR 0xc001102b write failed"
        elif [[ "$CPU_FAMILY" == "25" ]]; then
            # Family 25 includes both Zen3 and Zen4 - need model detection
            if [[ "$CPU_MODEL" -ge 96 && "$CPU_MODEL" -le 111 ]]; then
                echo -e "${GREEN}[INFO]${NOCOLOR} Detected Ryzen (Zen4)"
                # Zen4 MSR values - optimized for Raphael/Phoenix
                wrmsr -a 0xc0011020 0x4400000000000 2>/dev/null || echo -e "${YELLOW}[WARN]${NOCOLOR} MSR 0xc0011020 write failed"
                wrmsr -a 0xc0011021 0x4000000000040 2>/dev/null || echo -e "${YELLOW}[WARN]${NOCOLOR} MSR 0xc0011021 write failed"
                wrmsr -a 0xc0011022 0x8680000401570000 2>/dev/null || echo -e "${YELLOW}[WARN]${NOCOLOR} MSR 0xc0011022 write failed"
                wrmsr -a 0xc001102b 0x2040cc10 2>/dev/null || echo -e "${YELLOW}[WARN]${NOCOLOR} MSR 0xc001102b write failed"
            else
                echo -e "${GREEN}[INFO]${NOCOLOR} Detected Ryzen (Zen3)"
                # Zen3 MSR values - proven for Vermeer/Cezanne
                wrmsr -a 0xc0011020 0x4480000000000
                wrmsr -a 0xc0011021 0x1c000200000040
                wrmsr -a 0xc0011022 0xc000000401500000
                wrmsr -a 0xc001102b 0x2000cc14
            fi
        else
            echo -e "${GREEN}[INFO]${NOCOLOR} Detected Ryzen (Zen1/Zen2)"
            # Zen1/Zen2 MSR values - for Summit Ridge/Pinnacle Ridge/Matisse
            wrmsr -a 0xc0011020 0
            wrmsr -a 0xc0011021 0x40
            wrmsr -a 0xc0011022 0x1510000
            wrmsr -a 0xc001102b 0x2000cc16
        fi
        echo -e "${GREEN}[INFO]${NOCOLOR} AMD Ryzen RandomX boost applied"
    fi
}

huge_pages_set(){
    echo $1 > /proc/sys/vm/nr_hugepages
}

huge_cur_print(){
    echo -n -e "HugePages 2Mb (free/total): "
    if [[ $HPAGES_FREE -eq $HPAGES_CUR && $HPAGES_CUR -gt 0 ]]; then
        echo -e -n "${GREEN}"
    elif [[ $HPAGES_FREE -lt $HPAGES_CUR && $HPAGES_CUR -gt 0 ]]; then
        echo -e -n "${YELLOW}"
    else
        echo -e -n "${RED}"
    fi
    echo -n -e "$HPAGES_FREE${NOCOLOR} / "
    if [[ $HPAGES_CUR -gt 0 ]]; then
        echo -n -e "${GREEN}"
    else
        echo -n -e "${RED}"
    fi
    echo -e "$HPAGES_CUR${NOCOLOR}"
}

huge_cpu_print(){
    local AES=`lscpu | grep -c "aes"`
    echo -n -e "Your system has ${YELLOW}$CPU_NUMA${NOCOLOR} NUMA node(s) and ${YELLOW}$CPU_CORE${NOCOLOR} CPUs cores with "
    [[ $AES == 1 ]] && echo -e "${GREEN}AES support${NOCOLOR}" || echo -e "${RED}no AES support${NOCOLOR}"
}

huge_max_print(){
    echo -e "HugePages 2Mb (maximum RX): ${RED}$HPAGES_MAX${NOCOLOR}"
}

huge_rx_print(){
    echo -e "HugePages 2Mb (suggest RX): ${YELLOW}$HUGES_RX_CALC${NOCOLOR}"
    echo -n -e "HugePages 1Gb (RX-tunning): "
    [[ $MEM_AVAIL -gt $HPAGES_1GB_MEM ]] && echo -e "${GREEN}available${NOCOLOR}" || echo -e "${RED}not available${NOCOLOR}"
}

print_banner(){
    echo -e "${CYAN}HugePages Helper v${VERS}${NOCOLOR}"
}

print_help(){
    echo -e "${WHITE}This tool allows you to set the desired value for hugepages."
    echo -e "${LGRAY}The suggested value for CN-based algorithms is ${HUGES_CN}"
    echo -e "The suggested value for RandomX is $HUGES_RX_CALC"
    echo
    echo -e "${WHITE}Supported AMD architectures:${NOCOLOR}"
    echo -e "${LGRAY}  • Zen1/Zen2 (Ryzen 1000-3000 series) - Family 23"
    echo -e "  • Zen3 (Ryzen 5000 series) - Family 25, Models 0-47"
    echo -e "  • Zen4 (Ryzen 7000 series) - Family 25, Models 96-111"
    echo -e "  • Zen5 (Ryzen 9000 series) - Family 26"
    echo
    echo -e "${YELLOW}Usage: \n\t${CYAN}hugepages value${NOCOLOR}"
    echo
    echo -e "${YELLOW}Other examples of usage:${NOCOLOR}"
    echo -e "\t${CYAN}hugepages -r${NOCOLOR}   - ${LGRAY}restoring system defaults"
    echo -e "\t${CYAN}hugepages -rx${NOCOLOR}  - ${LGRAY}try set suggested value ($HUGES_RX_CALC) for RandomX"
    echo -e "\t${CYAN}hugepages -cn${NOCOLOR}  - ${LGRAY}try set suggested value ($HUGES_CN) for CN-based "
    echo -e "\t${CYAN}hugepages -erx${NOCOLOR} - ${LGRAY}try use $((CPU_NUMA*3))*1Gb HugePages + $CPU_CORE*2Mb HugePages for RandomX"
    echo -e "\t${CYAN}hugepages -mcN${NOCOLOR} - ${LGRAY}drop memory cache where N can be one of 1 (PageCache only), 2 (denries and inodes) or 3 (1+2)"
    echo -e "\t${CYAN}hugepages -h${NOCOLOR}   - ${LGRAY}this help "
}

huge_get_def(){
    local def=`cat /etc/sysctl.conf | grep -E "^vm.nr_hugepages" | tr '=' ' ' | awk '{ print $2 }'`
    [[ "$def" == "" ]] && echo 0 || echo $def
}

################################################################################
# MAIN SCRIPT BODY
################################################################################
cpu_boost=0

print_banner
[[ -z $HPAGES_REQ ]] && huge_cpu_print && huge_cur_print && huge_max_print && huge_rx_print && exit 0
case $1 in
    -h|--help)
        print_help
        exit 0
        ;;
    -r|--restore)
        huge_rx_1gb_pages_get
        if [[ $HPAGES_1GB_FREE -gt 0 ]]; then
           echo -e "${GREEN}[INFO]${NOCOLOR} Reset 1Gb HugePages - set to $((HPAGES_1GB_TOTAL-HPAGES_1GB_FREE))"
           huge_rx_1gb_pages_set $((HPAGES_1GB_TOTAL-HPAGES_1GB_FREE))
        fi
        HPAGES_REQ=`huge_get_def`
        echo -e "${GREEN}[INFO]${NOCOLOR} Try restore to system default HugePages value: ${YELLOW}$HPAGES_REQ${NOCOLOR}"
        ;;
    -mc1)
        memory_clear1
        exit 0
        ;;
    -mc2)
        memory_clear2
        exit 0
        ;;
    -mc3)
        memory_clear3
        exit 0
        ;;
    -erx)
        memory_clear3
        echo -e "${GREEN}[INFO]${NOCOLOR} Try use 1Gb HugePages for RandomX algorithm"
        cpu_boost=1
        huge_rx_1gb_pages_get
        if [[ $MEM_AVAIL -gt $HPAGES_1GB_MEM ]]; then
            huge_rx_1gb_pages_set
            huge_rx_1gb_pages_get
            if [[ $HPAGES_1GB_FREE -gt 0 ]]; then
                echo -e "${GREEN}[INFO]${NOCOLOR} Use ${YELLOW}${HPAGES_1GB_FREE}${NOCOLOR}x1Gb HugePages for RandomX algorithm"
                HPAGES_REQ=$(nproc)
            else
               echo -e "${YELLOW}[WARN]${NOCOLOR} Failed to allocate 1Gb HugePages, using 2Mb HugePages"
               HPAGES_REQ=$HUGES_RX_CALC
            fi
        else
            echo -e "${YELLOW}[WARN]${NOCOLOR} Not enough RAM available for 1Gb HugePages, using 2Mb HugePages"
            HPAGES_REQ=$HUGES_RX_CALC
        fi
        ;;
    -rx)
        cpu_boost=1
        HPAGES_REQ=$HUGES_RX_CALC
        echo -e "${GREEN}[INFO]${NOCOLOR} Try set HugePages to suggested for RandomX algorithm value: ${YELLOW}$HPAGES_REQ${NOCOLOR}"
        ;;
    -cn)
        HPAGES_REQ=128
        echo -e "${GREEN}[INFO]${NOCOLOR} Try set HugePages to suggested for CN-based algorithms value: ${YELLOW}$HPAGES_REQ${NOCOLOR}"
        ;;
     *)
        cpu_boost=1
        re='^[0-9]+$'
        if ! [[ $HPAGES_REQ =~ $re ]] ; then
            echo -e "${RED}[ERR] ${YELLOW}Invalid value given${NOCOLOR}" >&2; exit 1
        fi
esac

if [[ $cpu_boost -eq 1 ]]; then
    randomx_intel_boost
    randomx_ryzen_boost
fi

if [[ $HPAGES_REQ -le $HPAGES_MAX ]]; then
    huge_pages_set $HPAGES_REQ
else
    echo -e "${YELLOW}[WARN]${NOCOLOR} HugePages requested ($HPAGES_REQ) more then maximum suggested ($HPAGES_MAX)."
    huge_pages_set $HPAGES_MAX
fi

HPAGES_SET=`cat /proc/sys/vm/nr_hugepages`

echo -e "${GREEN}[INFO]${NOCOLOR} New HugePages (2Mb) value set: ${GREEN}$HPAGES_SET${NOCOLOR}"

exit 0
