#!/bin/bash

set -e

# TODO: update iwiki url
# Print error message and exit
msg_err() {
    echo "Error: $1" >&2
    echo "Please refer to https://iwiki.woa.com/p/4015781670" >&2
    exit 1
}

# 根据cpu个数计算安装nbs模块预期需要多少内存+agent需要多少内存
check_cpu_mem_usage() {
    local cgroup_mem_limit=$(echo "$1" | grep -oE '[0-9]+')

    # Check CPU usage
    cpu_count=$(nproc)
    expected_mem_usage=$(( cpu_count * 4 + cgroup_mem_limit * 1024 )) # 4MB per CPU

    mem_avail=$(grep MemAvailable /proc/meminfo | awk '{print $2;}')
    sleep 30
    mem_avail_again=$(grep MemAvailable /proc/meminfo | awk '{print $2;}')
    (( mem_avail_again < mem_avail )) && mem_avail=$mem_avail_again # 获取两次采样的空闲内存的较小值

    if (( mem_avail < expected_mem_usage )); then
        msg_err "The nbs system needs ${expected_mem_usage}MB but only gets available memory ${mem_avail}MB."
    fi
}

# 检查/proc/net/tcp的连接是否过多
# 如果要检查所有的连接数量，在本身/proc/nett/tcp中的连接数非常大量时可能会有较大耗时，且nbs_update未使用资源限制，不要做过重的检查操作
check_net_tcp_connections() {
    local max_connections=$1

    if [[ -f "/proc/net/tcp" ]]; then
        # 显示所有状态为01的连接详细信息
        local connections=$(awk 'NR>1 && $4=="01" {print $0}' /proc/net/tcp)
        local count=$(echo "$connections" | wc -l)

        if [ $count -gt $max_connections ]; then
            msg_err "The system has too many connections (>$max_connections links).Please check again."
        fi
    fi
}

# 检测是否应使用 MemoryMax 替代 MemoryLimit
# cgroupv2 场景，或 cgroupv1 + systemd >= 257 时使用 MemoryMax
get_memory_directive() {
    local cgroup_version=1
    # 检测 cgroup 版本：cgroup.controllers 存在说明是 cgroupv2
    if [[ -f /sys/fs/cgroup/cgroup.controllers ]]; then
        cgroup_version=2
    fi

    if [[ $cgroup_version -eq 2 ]]; then
        echo "MemoryMax"
        return
    fi

    # cgroupv1 场景：检查 systemd 版本
    local systemd_ver
    systemd_ver=$(systemctl --version 2>/dev/null | head -n1 | awk '{print $2}')
    if [[ -n "$systemd_ver" ]] && [[ "$systemd_ver" =~ ^[0-9]+$ ]] && (( systemd_ver >= 257 )); then
        echo "MemoryMax"
        return
    fi

    echo "MemoryLimit"
}

# 根据 parser.level 计算期望的 monitor_mode
# 3.0.0 起内核 monitor_mode 只保留两档：
#   0 (default): TS_SCM_TSTAMP_SND + TS_SCM_TSTAMP_ACK + TS_SCM_TSTAMP_RECV_TCP，迎合 send-ack 与 drv-l4 两段建接时延观测；
#   1 (all)    : 启用全部 tstype，覆盖 send 全链路与 recv 全链路。
get_expected_monitor_mode() {
    local parser_level="$1"

    case "$parser_level" in
        "all")
            echo "1"
            ;;
        "default"|*)
            echo "0"
            ;;
    esac
}



extract_release_number() {
    local kernel_version="$1"
    echo "$kernel_version" | sed -E 's/^[^-]*-([0-9]+).*$/\1/'
}

# 该函数的逻辑应该与internal/checker/ko_path.go中的getRecheckVer逻辑保持一致
get_recheck_ver() {
    local kernel_version="$1"

    # Check if the kernel version is in the /usr/lib/nbs directory
    # If so, return the kernel version directly, otherwise return rechecked version
    if [ -e /usr/lib/nbs/$kernel_version/nbs.ko ]; then
        echo "$kernel_version"
        return
    fi

    local release_number=$(extract_release_number "$kernel_version")

    if echo "$kernel_version" | grep -q "feat"; then
        echo "$kernel_version"
    elif echo "$kernel_version" | grep -q "-eks"; then
        echo "$kernel_version"
    elif echo "$kernel_version" | grep -q "\.eks"; then
        echo "$kernel_version" | sed -E 's/(.*\.eks)[^.]*$/\1/'
    elif echo "$kernel_version" | grep -qE '^6\.6\.110-42(\.([0-9]|[1-3][0-9]))?\.tl4\.x86_64|^6\.6\.119-49(\.([0-9]|1[0-9]|2[01]))?\.tl4\.x86_64'; then
        msg_err "kernel version $kernel_version has CVE vulnerabilities, please upgrade the kernel."
    elif echo "$kernel_version" | grep -qE '^6\.6\.110-42\.(4[0-9]|[5-9][0-9])\.tl4\.x86_64'; then
        echo "6.6.110-42.40.tl4.x86_64"
    elif echo "$kernel_version" | grep -qE '^6\.6\.119-49\.(2[2-9]|[3-9][0-9])\.tl4\.x86_64'; then
        echo "6.6.119-49.22.tl4.x86_64"
    elif echo "$kernel_version" | grep -q "^6\.6"; then
        echo "$kernel_version" | sed -E 's/^(.*-)([0-9]+)(\.[0-9]+)*(\.tl4\.x86_64)$/\1\2\4/'
    elif [[ "$release_number" -gt 1 ]]; then
        # return public kernels directly
        echo "$kernel_version"
    else
        echo "$kernel_version" | sed -E 's/^(.*-)([^.-]+)(\..*)?$/\1\2/'
    fi
}

# ====== Configuration ======
CONF_FILE="/etc/nbs_agent.conf"  # Path to the config file
TARGET_SERVICE_FILE="/usr/lib/systemd/system/nbs_agent.service"    # Default target service file path
TARGET_SERVICE_FILE_SRC="/opt/nbs_agent/misc/nbs_agent.service"

# Check for root privileges
if [[ $EUID -ne 0 ]]; then
    msg_err "Please run this script as root."
fi

# Check if config file exist
if [[ ! -f "$CONF_FILE" ]]; then
    msg_err "Config file $CONF_FILE not found. Please copy /opt/nbs_agent/misc/nbs_agent.conf to /etc/nbs_agent.conf and then run nbs_update again. Or reinstall nbs rpm directly."
fi


if [[ ! -f "$TARGET_SERVICE_FILE_SRC" ]]; then
    msg_err "Service template $TARGET_SERVICE_FILE_SRC not found!"
fi

# ====== Parse and export config ======
eval $(awk '
    /^resource_limit:/     { section="rl"; next }
    /^kernel:/             { section=""; next }
    /^parser:/             { section="p"; next }
    /^[a-z_]+:/ && !/^  /  { section=""; next }

    section=="rl" && /^  cpu_cores:/ {
        sub(/.*: */, ""); gsub(/"/, ""); gsub(/^[ \t]+|[ \t]+$/, ""); cpu_cores=$0; saw_cpu_cores=1
    }
    section=="rl" && /^  cpu_usage_limit:/ {
        sub(/.*: */, ""); gsub(/"/, ""); cpu_usage=$0
    }
    section=="rl" && /^  memory_limit:/ {
        sub(/.*: */, ""); gsub(/"/, ""); memory=$0
    }
    section=="p" && /^  level:/ {
        sub(/.*: */, ""); gsub(/"/, ""); parser_level=$0
    }

    END {
        # cpu_cores 为空表示不绑核，仅限制 CPU 利用率；未配置时默认绑核 "1"
        if (!saw_cpu_cores) cpu_cores="1"
        if (!cpu_usage) cpu_usage="50"
        if (!memory) memory="5G"
        if (!parser_level) parser_level="default"

        printf "CPU_CORES=\"%s\"\n", cpu_cores
        printf "CPU_USAGE_LIMIT=\"%s\"\n", cpu_usage
        printf "MEMORY_LIMIT=\"%s\"\n", memory
        printf "PARSER_LEVEL=\"%s\"\n", parser_level
    }
' "$CONF_FILE")

# ====== Parse kernel rules via python3 ======
eval $(python3 -c "
import yaml, sys
with open('$CONF_FILE') as f:
    conf = yaml.safe_load(f)
kw = conf.get('kernel', {}) if conf else {}
rules = kw.get('rules', [])
atomic_wrt = kw.get('atomic_wrt', 2)
print('RULE_COUNT={}'.format(len(rules)))
for i, r in enumerate(rules):
    for k in ['dir', 'family', 'src_ip', 'dst_ip', 'src_ports', 'dst_ports']:
        v = r.get(k, '') or ''
        print('RULE_{}_{}=\"{}\"'.format(i, k.upper(), v))
print('ATOMIC_WRT=\"{}\"'.format(atomic_wrt))
" 2>/dev/null) || {
    echo "Warning: Failed to parse kernel.rules with python3. Rules default to empty."
    RULE_COUNT=0
}

# ====== Parameter validation ======
# Get the number of available CPU cores on this machine
AVAILABLE_CORES=$(nproc)

# cpu_cores 为空：不绑核，仅用 CPUQuota 限制利用率
# cpu_cores 非空：CPU 核编号（systemd CPUAffinity，0-based），合法范围 0 .. nproc-1
if [[ -n "$CPU_CORES" ]]; then
    # 允许 0（0 号核合法）；不允许前导零，避免 "01" 之类
    if ! [[ "$CPU_CORES" =~ ^(0|[1-9][0-9]*)$ ]]; then
        msg_err "CPU_CORES must be a non-negative integer (0-based CPU core id), or empty to disable CPU affinity."
    fi

    # 核编号从 0 开始，最大合法值为 AVAILABLE_CORES-1
    if (( CPU_CORES >= AVAILABLE_CORES )); then
        msg_err "CPU_CORES ($CPU_CORES) is out of range; valid core ids are 0..$((AVAILABLE_CORES - 1))."
    fi
    echo "Info: CPU affinity enabled, CPUAffinity=$CPU_CORES"
else
    echo "Info: cpu_cores is empty, CPU affinity disabled (CPUQuota only)."
fi

# Check if CPU_USAGE_LIMIT is a positive integer between 1 and 100
if ! [[ "$CPU_USAGE_LIMIT" =~ ^[1-9][0-9]*$ ]]; then
    msg_err "CPU_USAGE_LIMIT must be a positive integer."
fi

if (( CPU_USAGE_LIMIT < 1 || CPU_USAGE_LIMIT > 100 )); then
    msg_err "CPU_USAGE_LIMIT ($CPU_USAGE_LIMIT) must be between 1 and 100."
fi

# Check if MEMORY_LIMIT is larger or equal than 1G.
if [[ "$MEMORY_LIMIT" =~ ^([0-9]+)G$ ]]; then
    number="${BASH_REMATCH[1]}"
    if (( number < 1 )); then
        msg_err "MEMORY_LIMIT ($MEMORY_LIMIT) must be larger or equal than 1G."
    fi
else
    msg_err "MEMORY_LIMIT ($MEMORY_LIMIT) should be consist of a number and G, like 5G."
fi

check_cpu_mem_usage "$MEMORY_LIMIT"
# we set the max connections to 10000 for this version
# check_net_tcp_connections 10000

# 校验 ATOMIC_WRT 值（只允许 0/1/2）
if ! [[ "$ATOMIC_WRT" =~ ^[0-2]$ ]]; then
    echo "Warning: Invalid atomic_wrt=$ATOMIC_WRT, defaulting to 2."
    ATOMIC_WRT=2
fi

# ====== Check and update kernel module monitor_mode ======
kernel_full_ver=$(uname -r)
kernel_recheck_ver=$(get_recheck_ver "$kernel_full_ver")
echo "Info: Detected kernel version: $kernel_recheck_ver (full: $kernel_full_ver)"

NBS_KO="/usr/lib/nbs/$kernel_recheck_ver/nbs.ko"
NBS_MODULE_PATH="/sys/module/nbs"
MONITOR_MODE_FILE="/sys/module/nbs/parameters/monitor_mode"

insmod_nbs_ko() {
    insmod "$NBS_KO" monitor_mode="$EXPECTED_MONITOR_MODE"
}

insmod_nbs_ko_log_suffix() {
    echo "monitor_mode=$EXPECTED_MONITOR_MODE"
}

# 计算期望的 monitor_mode
EXPECTED_MONITOR_MODE=$(get_expected_monitor_mode "$PARSER_LEVEL")
echo "Info: Expected monitor_mode=$EXPECTED_MONITOR_MODE (level=$PARSER_LEVEL)"

if [ ! -e "$NBS_MODULE_PATH" ]; then
    # ko内核模块未安装，直接安装ko，并带上 monitor_mode 参数
    echo "Info: nbs ko has not been loaded. We will load it with $(insmod_nbs_ko_log_suffix)"
    if [[ -f "$NBS_KO" ]]; then
        if ! insmod_nbs_ko; then
            msg_err "Failed to load $NBS_KO with $(insmod_nbs_ko_log_suffix)"
        fi
        echo "Info: nbs kernel module loaded with $(insmod_nbs_ko_log_suffix)"
    else
        msg_err "The nbs kernel module not found: $NBS_KO"
    fi
elif [[ -f "$MONITOR_MODE_FILE" ]]; then
    # ko内核模块已经安装，检查模块的 monitor_mode
    CURRENT_MONITOR_MODE=$(cat "$MONITOR_MODE_FILE")
    echo "Info: Current kernel monitor_mode=$CURRENT_MONITOR_MODE"
    
    # 验证 CURRENT_MONITOR_MODE 是否为有效数字（0-1）
    if ! [[ "$CURRENT_MONITOR_MODE" =~ ^[01]$ ]]; then
        msg_err "Invalid monitor_mode value '$CURRENT_MONITOR_MODE' read from $MONITOR_MODE_FILE. Expected 0 or 1."
    fi
    
    if [[ "$CURRENT_MONITOR_MODE" != "$EXPECTED_MONITOR_MODE" ]]; then
        echo "Info: monitor_mode mismatch, reloading nbs kernel module..."
        
        # 先停止 nbs_agent 服务（如果正在运行）
        if systemctl is-active --quiet nbs_agent 2>/dev/null; then
            echo "Info: Stopping nbs_agent service before reloading kernel module..."
            systemctl stop nbs_agent || true
        fi
        
        # 卸载旧的内核模块
        if ! rmmod nbs; then
            msg_err "Failed to unload nbs kernel module."
        fi
        echo "Info: nbs kernel module unloaded."
        
        # 加载新的内核模块，带上 monitor_mode 参数
        if [[ -f "$NBS_KO" ]]; then
            if ! insmod_nbs_ko; then
                msg_err "Failed to load $NBS_KO with $(insmod_nbs_ko_log_suffix)"
            fi
            echo "Info: nbs kernel module loaded with $(insmod_nbs_ko_log_suffix)"
        else
            msg_err "The nbs kernel module not found: $NBS_KO"
        fi
    else
        echo "Info: monitor_mode already matches, no need to reload kernel module."
    fi
else
    # 内核模块不支持 monitor_mode 参数，默认只支持 monitor_mode=0 (level=default)
    echo "Warning: $MONITOR_MODE_FILE not found, kernel module only supports monitor_mode=0 (level=default)."

    # 旧内核模块只支持 default，all 不兼容
    if [[ "$PARSER_LEVEL" != "default" ]]; then
        msg_err "Kernel module does not support monitor_mode parameter, only level=default is supported. Current config (level=$PARSER_LEVEL) is not compatible."
    fi

    echo "Info: Agent config is compatible with kernel module (level=default)."
fi

# ====== Set atomic_wrt parameter ======
ATOMIC_WRT_FILE="/sys/module/nbs/parameters/atomic_wrt"
if [[ -f "$ATOMIC_WRT_FILE" ]]; then
    echo "$ATOMIC_WRT" > "$ATOMIC_WRT_FILE" 2>/dev/null || echo "Warning: Failed to set atomic_wrt=$ATOMIC_WRT"
    echo "Info: atomic_wrt set to $ATOMIC_WRT"
else
    echo "Info: atomic_wrt parameter not supported by kernel module, skipping."
fi

# ====== Build new service file ======
cp "$TARGET_SERVICE_FILE_SRC" "$TARGET_SERVICE_FILE"

# Replace template variables with actual values
# cpu_cores 为空时删除绑核行，只保留 CPUQuota
if [[ -n "$CPU_CORES" ]]; then
    sed -i "s|{{CPU_AFFINITY_LINE}}|CPUAffinity=$CPU_CORES|g" "$TARGET_SERVICE_FILE"
else
    sed -i '/{{CPU_AFFINITY_LINE}}/d' "$TARGET_SERVICE_FILE"
fi
sed -i "s|{{CPU_USAGE_LIMIT}}|${CPU_USAGE_LIMIT}%|g" "$TARGET_SERVICE_FILE"
# 根据 cgroup 版本和 systemd 版本决定使用 MemoryLimit 还是 MemoryMax
MEMORY_DIRECTIVE=$(get_memory_directive)
echo "Info: Using $MEMORY_DIRECTIVE for memory resource control."
sed -i "s|{{MEMORY_DIRECTIVE}}|$MEMORY_DIRECTIVE|g" "$TARGET_SERVICE_FILE"
sed -i "s|{{MEMORY_LIMIT}}|$MEMORY_LIMIT|g" "$TARGET_SERVICE_FILE"

chmod 644 "$TARGET_SERVICE_FILE"

# ====== Force tagent regenerate machine.ip ======
rm -f /usr/local/tagent/var/machine.ip || true

# ====== Reload systemd and start service ======
systemctl daemon-reload

# Check if the service exists
if systemctl list-unit-files | grep -q '^nbs_agent\.service'; then
    if systemctl is-active --quiet nbs_agent; then
        systemctl stop nbs_agent || msg_err "Failed to stop nbs_agent service."
        systemctl start nbs_agent || msg_err "Failed to start nbs_agent service."
    else
        systemctl start nbs_agent || msg_err "Failed to start nbs_agent service."
    fi
else
    # 不打印 warning，只尝试启动
    systemctl start nbs_agent || msg_err "Failed to start nbs_agent service."
fi

# ====== Set traffic filter rules via nbsctl ======
NBSCTL_BIN="/usr/bin/nbsctl"
if [[ ! -x "$NBSCTL_BIN" ]]; then
    msg_err "$NBSCTL_BIN not found or not executable. New NBS module required."
fi

# 清空现有规则
if [ -e /dev/nbs ]; then
    "$NBSCTL_BIN" flush
    echo "Info: nbsctl rules flushed."
else
    msg_err "nbs ko is not installed successfully. There is no /dev/nbs."
fi

# 逐组下发过滤规则
has_rule=false
for ((i=0; i<RULE_COUNT; i++)); do
    eval "dir=\"\${RULE_${i}_DIR:-}\""
    eval "fam=\"\${RULE_${i}_FAMILY:-}\""
    eval "sip=\"\${RULE_${i}_SRC_IP:-}\""
    eval "dip=\"\${RULE_${i}_DST_IP:-}\""
    eval "sp=\"\${RULE_${i}_SRC_PORTS:-}\""
    eval "dp=\"\${RULE_${i}_DST_PORTS:-}\""

    args=""
    [[ -n "$dir" ]] && args="$args dir=$dir"
    [[ -n "$fam" ]] && args="$args family=$fam"
    [[ -n "$sip" ]] && args="$args src=$sip"
    [[ -n "$dip" ]] && args="$args dst=$dip"
    [[ -n "$sp" ]]  && args="$args sport=$sp"
    [[ -n "$dp" ]]  && args="$args dport=$dp"

    if [[ -n "$args" ]]; then
        "$NBSCTL_BIN" add $args
        echo "Info: nbsctl add$args"
        has_rule=true
    fi
done

if ! $has_rule; then
    echo "Warning: No filter rules specified. No flow tracing will be performed."
fi

# enable nbs（规则下发后再启用，保证新建连接以及存量都被正确标记）
sysctl -w net.nbs.enable_nbs=2

# ====== Run nbs_check ======
NBS_CHECK_BIN="/usr/bin/nbs_check"
if [[ -x "$NBS_CHECK_BIN" ]]; then
    echo "Info: Running nbs_check (skipping result checks since agent just started)..."
    "$NBS_CHECK_BIN" -skip=result || true
else
    echo "Warning: $NBS_CHECK_BIN not found, skipping self-check."
fi

exit 0
