#!/bin/bash
set -euo pipefail

# =========================================================
# TuneTCP - TCP参数优化工具
# 作者: Yuju
# 功能: 自动检测并优化TCP参数，启用BBR拥塞控制
# =========================================================

# 颜色定义
readonly RED='\033[1;31m'
readonly GREEN='\033[1;32m'
readonly YELLOW='\033[1;33m'
readonly BLUE='\033[1;34m'
readonly PURPLE='\033[1;35m'
readonly CYAN='\033[1;36m'
readonly WHITE='\033[1;37m'
readonly RESET='\033[0m'

# 日志目录
readonly LOG_FILE="/var/log/tunetcp.log"

# 输出函数
log() { echo "$(date '+%Y-%m-%d %H:%M:%S') - $*" >> "$LOG_FILE"; }
info() { echo -e "${BLUE}[INFO]${RESET} $*" >&2; log "INFO: $*"; }
success() { echo -e "${GREEN}[SUCCESS]${RESET} $*" >&2; log "SUCCESS: $*"; }
warning() { echo -e "${YELLOW}[WARNING]${RESET} $*" >&2; log "WARNING: $*"; }
error() { echo -e "${RED}[ERROR]${RESET} $*" >&2; log "ERROR: $*"; exit 1; }

# 检查root权限
check_root() {
    [[ $EUID -eq 0 ]] || error "请以root权限运行此脚本"
}

# 检测操作系统类型
check_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        echo "${ID:-unknown}"
    elif command -v lsb_release >/dev/null 2>&1; then
        lsb_release -si | tr '[:upper:]' '[:lower:]'
    else
        echo "unknown"
    fi
}

# 检查内核版本
check_kernel_version() {
    local kernel_version
    kernel_version=$(uname -r | cut -d. -f1-2)
    local major minor
    major=$(echo "$kernel_version" | cut -d. -f1)
    minor=$(echo "$kernel_version" | cut -d. -f2)
    
    if [[ $major -lt 4 ]] || [[ $major -eq 4 && $minor -lt 9 ]]; then
        warning "内核版本 ${kernel_version} 可能不完全支持BBR，建议升级到4.9+"
        return 1
    fi
    success "内核版本 ${kernel_version} 支持BBR"
    return 0
}

# 创建备份目录
create_backup_dir() {
    local timestamp=$(date +%Y%m%d-%H%M%S)
    local backup_base="/etc/tunetcp-backup"
    local backup_dir="$backup_base/$timestamp"
    
    # 创建基础目录
    mkdir -p "$backup_base"
    
    # 返回备份路径，但暂不创建具体目录
    echo "$backup_dir"
}
# 添加实际创建备份目录的函数
ensure_backup_dir() {
    local backup_dir="$1"
    if [[ ! -d "$backup_dir" ]]; then
        mkdir -p "$backup_dir"
        info "创建备份目录: $backup_dir"
    fi
}
# 内存检测
check_mem() {
    local mem_bytes
    if [[ -r /proc/meminfo ]]; then
        mem_bytes=$(awk '/^MemTotal:/ {print $2 * 1024}' /proc/meminfo)
    elif command -v free >/dev/null 2>&1; then
        mem_bytes=$(free -b | awk 'NR==2 {print $2}')
    elif [[ -r /proc/meminfo ]]; then
        mem_bytes=$(grep '^MemTotal:' /proc/meminfo | awk '{print $2 * 1024}')
    else
        error "无法获取内存信息：系统不支持"
        return 1
    fi
    if [[ ! "$mem_bytes" =~ ^[0-9]+$ ]] || (( mem_bytes < 134217728 )); then  # 最小128MB
        error "获取的内存信息无效: ${mem_bytes:-"空值"} bytes"
        return 1
    fi
    printf "%.2f" $((mem_bytes * 100 / 1073741824))e-2
}

rtt_test() {
    local ping_target=""
    local ping_desc=""
    local -a dns_servers=("1.1.1.1" "8.8.8.8" "223.5.5.5" "114.114.114.114")
    local timeout_duration=5
    local ping_count=3
    local default_rtt=50
    validate_ip() {
        local ip="$1"
        if [[ $ip =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; then
            local IFS='.'
            local -a ip_parts=($ip)
            for part in "${ip_parts[@]}"; do
                [[ $part -le 255 ]] || return 1
            done
            return 0
        fi
        return 1
    }
    ping_target_host() {
        local target="$1"
        local desc="$2"
        
        info "正在测试到 ${desc} 的延迟..."
        local ping_output
        if ping_output=$(timeout "$timeout_duration" ping -c "$ping_count" -W 2 -i 0.5 "$target" 2>/dev/null); then
            local rtt
            if command -v bc >/dev/null 2>&1; then
                rtt=$(echo "$ping_output" | awk -F'[=/]' '/rtt|round-trip/ {print $5}' | head -1)
            else
                rtt=$(echo "$ping_output" | awk -F'[=/]' '/rtt|round-trip/ {printf "%.0f", $5}' | head -1)
            fi
            
            if [[ "$rtt" =~ ^[0-9]+(\.[0-9]+)?$ ]] && (( $(echo "$rtt > 0" | bc -l 2>/dev/null || echo "1") )); then
                success "检测到平均RTT: ${rtt} ms"
                printf "%.0f" "$rtt"
                return 0
            fi
        fi
        
        warning "Ping ${target} 失败或结果无效"
        return 1
    }
    test_multiple_targets() {
        local -a targets=("$@")
        local rtt_sum=0
        local valid_tests=0
        local -a results=()
        info "使用多个DNS服务器测试网络延迟..."
        
        for target in "${targets[@]}"; do
            local rtt
            if rtt=$(timeout "$timeout_duration" ping -c "$ping_count" -W 2 -i 0.5 "$target" 2>/dev/null | \
                     awk -F'[=/]' '/rtt|round-trip/ {printf "%.1f", $5; exit}'); then
                
                if [[ "$rtt" =~ ^[0-9]+(\.[0-9]+)?$ ]] && (( $(echo "$rtt > 0" | bc -l 2>/dev/null || echo "1") )); then
                    if command -v bc >/dev/null 2>&1; then
                        rtt_sum=$(echo "$rtt_sum + $rtt" | bc -l)
                    else
                        rtt_sum=$(awk -v sum="$rtt_sum" -v rtt="$rtt" 'BEGIN {print sum + rtt}')
                    fi
                    ((valid_tests++))
                    results+=("DNS ${target}: ${rtt} ms")
                    info "DNS ${target}: ${rtt} ms"
                else
                    warning "DNS ${target} 测试结果解析失败"
                fi
            else
                warning "DNS ${target} 连接超时或失败"
            fi
        done
        
        if [[ $valid_tests -gt 0 ]]; then
            local avg_rtt
            if command -v bc >/dev/null 2>&1; then
                avg_rtt=$(echo "scale=0; $rtt_sum / $valid_tests" | bc -l)
            else
                avg_rtt=$(awk -v sum="$rtt_sum" -v count="$valid_tests" 'BEGIN {printf "%.0f", sum / count}')
            fi
            
            success "平均RTT (${valid_tests}个有效测试): ${avg_rtt} ms"
            echo "$avg_rtt"
            return 0
        fi
        
        return 1
    }
    if [[ -n "${SSH_CONNECTION:-}" ]]; then
        ping_target=$(echo "$SSH_CONNECTION" | awk '{print $1}')
        if validate_ip "$ping_target"; then
            ping_desc="SSH 客户端 ${ping_target}"
            info "检测到SSH连接，客户端IP: ${ping_target}"
        else
            warning "SSH连接IP格式无效: $ping_target"
            ping_target=""
        fi
    fi
    if [[ -z "$ping_target" ]]; then
        info "未检测到有效SSH连接，请提供测试目标"
        local client_ip
        read -r -p "请输入客户机IP (回车使用多个公共DNS测试): " client_ip
        
        if [[ -n "$client_ip" ]]; then
            if validate_ip "$client_ip"; then
                ping_target="$client_ip"
                ping_desc="客户机IP ${ping_target}"
            else
                warning "输入的IP地址格式无效: $client_ip"
                ping_target=""
            fi
        fi
    fi
    if [[ -n "$ping_target" ]]; then
        if ping_target_host "$ping_target" "$ping_desc"; then
            return 0
        fi
        warning "目标 ${ping_target} 测试失败，转为使用DNS服务器测试"
    fi
    if test_multiple_targets "${dns_servers[@]}"; then
        return 0
    fi
    warning "所有网络测试失败，使用默认值 ${default_rtt}ms"
    echo "$default_rtt"
    return 1
}


# 智能带宽检测
detect_bandwidth() {
    local bandwidth=1000  # 默认带宽值 (Mbps)
    local default_iface
    default_iface=$(ip -o -4 route show to default 2>/dev/null | awk '{print $5}' | head -1)
    
    if [[ -n "$default_iface" && -r "/sys/class/net/$default_iface/speed" ]]; then
        local speed
        speed=$(cat "/sys/class/net/$default_iface/speed" 2>/dev/null)
        
        if [[ "$speed" =~ ^[0-9]+$ && $speed -gt 0 && $speed -ne 4294967295 ]]; then
            bandwidth=$speed
            info "检测到网卡 ${default_iface} 速度: ${speed} Mbps"
        else
            info "网卡 ${default_iface} 速度检测失败 (值: ${speed:-未知})，使用默认值 ${bandwidth} Mbps"
        fi
    else
        info "无法读取网卡速度信息，使用默认值 ${bandwidth} Mbps"
    fi
    
    echo "$bandwidth"
}

# 验证参数有效性
validate_params() {
    local mem_g="$1"
    local bw_mbps="$2" 
    local rtt_ms="$3"
    
    # 内存检查 (0.5GB - 1024GB)
    if ! awk -v m="$mem_g" 'BEGIN {exit (m >= 0.5 && m <= 1024) ? 0 : 1}'; then
        error "内存参数无效: ${mem_g} GiB (应在 0.5-1024 范围内)"
    fi
    
    # 带宽检查 (1Mbps - 100Gbps)
    if ! awk -v b="$bw_mbps" 'BEGIN {exit (b >= 1 && b <= 100000) ? 0 : 1}'; then
        error "带宽参数无效: ${bw_mbps} Mbps (应在 1-100000 范围内)"
    fi
    
    # RTT检查 (1ms - 1000ms)
    if ! awk -v r="$rtt_ms" 'BEGIN {exit (r >= 1 && r <= 1000) ? 0 : 1}'; then
        error "RTT参数无效: ${rtt_ms} ms (应在 1-1000 范围内)"
    fi
}

# 高级缓冲区计算
calculate_buffers() {
    local mem_g="$1"
    local bw_mbps="$2"
    local rtt_ms="$3"
    
    # BDP计算: 带宽 * RTT
    local bdp_bytes
    bdp_bytes=$(awk -v bw="$bw_mbps" -v rtt="$rtt_ms" 'BEGIN{ printf "%.0f", bw * 125 * rtt }')
    
    # 内存限制: 总内存的3%
    local mem_bytes ram_limit
    mem_bytes=$(awk -v g="$mem_g" 'BEGIN{ printf "%.0f", g * 1024^3 }')
    ram_limit=$(awk -v m="$mem_bytes" 'BEGIN{ printf "%.0f", m * 0.03 }')
    
    # 系统限制: 64MB (避免过度消耗内存)
    local sys_limit=$((64 * 1024 * 1024))
    
    # 最终缓冲区大小: min(2*BDP, 3%RAM, 64MB)
    local max_buffer
    max_buffer=$(awk -v bdp="$bdp_bytes" -v ram="$ram_limit" -v sys="$sys_limit" '
        BEGIN {
            max_buf = bdp * 2
            if (ram < max_buf) max_buf = ram
            if (sys < max_buf) max_buf = sys
            printf "%.0f", max_buf
        }')
    
    # 缓冲区分档 (便于管理)
    local max_mb buffer_mb
    max_mb=$((max_buffer / 1024 / 1024))
    
    if [[ $max_mb -ge 64 ]]; then buffer_mb=64
    elif [[ $max_mb -ge 32 ]]; then buffer_mb=32
    elif [[ $max_mb -ge 16 ]]; then buffer_mb=16
    elif [[ $max_mb -ge 8 ]]; then buffer_mb=8
    elif [[ $max_mb -ge 4 ]]; then buffer_mb=4
    else buffer_mb=4
    fi
    
    local final_bytes=$((buffer_mb * 1024 * 1024))
    
    # 返回计算结果
    echo "$bdp_bytes $final_bytes $buffer_mb"
}

# 优化的冲突处理
handle_conflicts() {
    local backup_dir="$1"
    local key_regex='^[[:space:]]*net\.(core\.(default_qdisc|rmem_max|wmem_max|rmem_default|wmem_default)|ipv4\.tcp_(rmem|wmem|congestion_control))[[:space:]]*='
    local backup_created=false
    
    info "处理配置冲突..."
    
    # 备份并处理 /etc/sysctl.conf
    if [[ -f "/etc/sysctl.conf" ]]; then
        if grep -Eq "$key_regex" "/etc/sysctl.conf"; then
            ensure_backup_dir "$backup_dir"
            backup_created=true
            cp "/etc/sysctl.conf" "$backup_dir/sysctl.conf.orig"
            info "已备份 /etc/sysctl.conf"
            # ... 其余处理逻辑保持不变
        fi
    fi
    
    # 处理 /etc/sysctl.d/ 中的其他文件
    if [[ -d "/etc/sysctl.d" ]]; then
        local target_file="/etc/sysctl.d/yuju-net-bbr-fq.conf"
        for conf_file in /etc/sysctl.d/*.conf; do
            [[ -f "$conf_file" ]] || continue
            [[ "$(readlink -f "$conf_file" 2>/dev/null)" == "$(readlink -f "$target_file" 2>/dev/null)" ]] && continue
            
            if grep -Eq "$key_regex" "$conf_file"; then
                if [[ "$backup_created" == false ]]; then
                    ensure_backup_dir "$backup_dir"
                    backup_created=true
                fi
                local basename_file=$(basename "$conf_file")
                cp "$conf_file" "$backup_dir/$basename_file"
                if mv "$conf_file" "$conf_file.disabled-by-tunetcp"; then
                    info "已禁用冲突文件: $conf_file"
                else
                    warning "无法禁用文件: $conf_file"
                fi
            fi
        done
    fi
    
    # 如果没有实际备份，返回空字符串
    if [[ "$backup_created" == false ]]; then
        echo ""
    else
        echo "$backup_dir"
    fi
}

# 应用网络优化配置
apply_optimization() {
    local buffer_bytes="$1"
    local buffer_mb="$2"
    local mem_g="$3"
    local bw_mbps="$4"
    local rtt_ms="$5"
    local bdp_bytes="$6"
    
    local config_file="/etc/sysctl.d/yuju-net-bbr-fq.conf"
    
    # 根据缓冲区大小设置默认值
    local def_rmem def_wmem
    if [[ $buffer_mb -ge 32 ]]; then
        def_rmem=262144; def_wmem=524288
    elif [[ $buffer_mb -ge 8 ]]; then
        def_rmem=131072; def_wmem=262144
    else
        def_rmem=131072; def_wmem=131072
    fi
    
    # TCP缓冲区参数
    local tcp_rmem_min=4096 tcp_rmem_def=87380 tcp_rmem_max=$buffer_bytes
    local tcp_wmem_min=4096 tcp_wmem_def=65536 tcp_wmem_max=$buffer_bytes
    
    info "生成优化配置..."
    
    # 生成配置文件
    cat > "$config_file" <<EOF
# TuneTCP Auto-generated Configuration
# Generated: $(date)
# Parameters: MEM=${mem_g}GiB, BW=${bw_mbps}Mbps, RTT=${rtt_ms}ms
# BDP: ${bdp_bytes} bytes (~$(awk -v b="$bdp_bytes" 'BEGIN{ printf "%.2f", b/1024/1024 }') MB)
# Buffer Size: ${buffer_mb} MB

# BBR拥塞控制
net.core.default_qdisc = fq
net.ipv4.tcp_congestion_control = bbr

# 核心网络缓冲区
net.core.rmem_default = ${def_rmem}
net.core.wmem_default = ${def_wmem}
net.core.rmem_max = ${buffer_bytes}
net.core.wmem_max = ${buffer_bytes}

# TCP缓冲区配置
net.ipv4.tcp_rmem = ${tcp_rmem_min} ${tcp_rmem_def} ${tcp_rmem_max}
net.ipv4.tcp_wmem = ${tcp_wmem_min} ${tcp_wmem_def} ${tcp_wmem_max}

# TCP优化选项
net.ipv4.tcp_mtu_probing = 1
net.ipv4.tcp_fastopen = 3
net.ipv4.tcp_slow_start_after_idle = 0
net.ipv4.tcp_notsent_lowat = 16384

# 连接跟踪优化 (适用于高并发)
net.netfilter.nf_conntrack_max = 1048576
net.netfilter.nf_conntrack_tcp_timeout_established = 1200

# 其他网络优化
net.core.netdev_max_backlog = 5000
net.ipv4.tcp_max_syn_backlog = 8192
net.core.somaxconn = 32768
EOF

    success "配置文件已生成: $config_file"
    
    # 应用配置
    info "应用sysctl配置..."
    if sysctl --system >/dev/null 2>&1; then
        success "sysctl配置应用成功"
    else
        error "sysctl配置应用失败"
    fi
    
    # 配置网络接口队列
    local default_iface
    default_iface=$(ip -o -4 route show to default 2>/dev/null | awk '{print $5}' | head -1)
    
    if command -v tc >/dev/null 2>&1 && [[ -n "$default_iface" ]]; then
        info "配置接口 ${default_iface} 的队列调度..."
        if tc qdisc replace dev "$default_iface" root fq 2>/dev/null; then
            success "接口队列配置成功"
        else
            warning "接口队列配置失败 (可能需要重启生效)"
        fi
    fi
}

# 显示优化结果
show_results() {
    local mem_g="$1"
    local bw_mbps="$2"
    local rtt_ms="$3"
    local bdp_bytes="$4"
    local buffer_mb="$5"
    
    local bdp_mb
    bdp_mb=$(awk -v b="$bdp_bytes" 'BEGIN{ printf "%.2f", b/1024/1024 }')
    
    echo
    echo "==== [ TuneTCP 优化结果 ] ===="
    echo
    
    # 参数摘要
    echo -e "${GREEN}[+] 优化参数摘要${RESET}"
    printf "    - %-15s : %s\n" "内存大小" "${mem_g} GiB"
    printf "    - %-15s : %s\n" "网络带宽" "${bw_mbps} Mbps"
    printf "    - %-15s : %s\n" "网络延迟" "${rtt_ms} ms"
    printf "    - %-15s : %s\n" "计算BDP" "${bdp_mb} MB"
    printf "    - %-15s : %s\n" "缓冲区大小" "${buffer_mb} MB"
    echo
    
    # 当前配置验证
    echo -e "${GREEN}[+] 当前配置验证${RESET}"
    printf "    - %-25s : %s\n" "TCP拥塞控制" "$(sysctl -n net.ipv4.tcp_congestion_control 2>/dev/null || echo '未知')"
    printf "    - %-25s : %s\n" "默认队列调度" "$(sysctl -n net.core.default_qdisc 2>/dev/null || echo '未知')"
    printf "    - %-25s : %s\n" "最大接收缓冲区" "$(sysctl -n net.core.rmem_max 2>/dev/null || echo '未知')"
    printf "    - %-25s : %s\n" "最大发送缓冲区" "$(sysctl -n net.core.wmem_max 2>/dev/null || echo '未知')"
    echo
    
    # 网络接口状态
    local default_iface
    default_iface=$(ip -o -4 route show to default 2>/dev/null | awk '{print $5}' | head -1)
    if command -v tc >/dev/null 2>&1 && [[ -n "$default_iface" ]]; then
        echo -e "${GREEN}[+] 网络接口状态${RESET}"
        printf "    - %-15s : %s\n" "默认接口" "$default_iface"
        local queue_type
        queue_type=$(tc qdisc show dev "$default_iface" 2>/dev/null | head -1 | awk '{print $2}' || echo "未知")
        printf "    - %-15s : %s\n" "队列类型" "$queue_type"
        echo
    fi
    
    echo -e "${CYAN}备份文件位置: ${backup_dir}${RESET}"
    echo -e "${CYAN}配置文件位置: /etc/sysctl.d/yuju-net-bbr-fq.conf${RESET}"
    echo -e "${CYAN}日志文件位置: $LOG_FILE${RESET}"
    echo "=================================="
}

# 主函数
main() {
    echo -e "${CYAN}"
    echo "╔══════════════════════════════════════════════════╗"
    echo "║            智能TCP网络参数优化工具               ║"
    echo "╚══════════════════════════════════════════════════╝"
    echo -e "${RESET}"
    
    # 权限检查
    check_root
    
    # 系统检查
    local os_type
    os_type=$(check_os)
    info "检测到操作系统: $os_type"
    
    check_kernel_version || warning "建议升级内核以获得更好的BBR支持"
    
    # 创建备份目录
    local backup_dir
    backup_dir=$(create_backup_dir)
    info "备份目录: $backup_dir"
    
    # 自动检测参数
    info "正在自动检测系统参数..."
    local mem_g bw_mbps rtt_ms
    mem_g=$(check_mem)
    bw_mbps=$(detect_bandwidth)
    rtt_ms=$(rtt_test)
    
    success "自动检测完成"
    
    # 参数确认循环
    while true; do
        clear
        echo -e "${CYAN}╔══════════════════════════════════════════════════════════════════╗${RESET}"
        echo -e "${WHITE}                        网络参数配置确认                          ${RESET}"
        echo -e "${CYAN}╠══════════════════════════════════════════════════════════════════╣${RESET}"
        echo -e "${WHITE} 1.${RESET} 内存大小     : ${GREEN}${mem_g} GiB${RESET} ${WHITE}(系统自动检测)${RESET}"
        echo -e "${WHITE} 2.${RESET} 网络带宽     : ${GREEN}${bw_mbps} Mbps${RESET} ${WHITE}(智能检测)${RESET}"
        echo -e "${WHITE} 3.${RESET} 网络延迟     : ${GREEN}${rtt_ms} ms${RESET} ${WHITE}(网络测试)${RESET}"
        echo -e "${CYAN}╚══════════════════════════════════════════════════════════════════╝${RESET}"
        
        local preview_result
        preview_result=$(calculate_buffers "$mem_g" "$bw_mbps" "$rtt_ms")
        local bdp_bytes buffer_bytes buffer_mb
        read -r bdp_bytes buffer_bytes buffer_mb <<< "$preview_result"
        local bdp_mb
        bdp_mb=$(awk -v b="$bdp_bytes" 'BEGIN{ printf "%.2f", b/1024/1024 }')
        echo
        
        read -r -p "按 [Enter] 应用优化, 输入 [1-3] 修改参数, [q] 退出: " choice
        
        case "$choice" in
            "")
                validate_params "$mem_g" "$bw_mbps" "$rtt_ms"
                info "参数验证通过，开始优化..."
                break
                ;;
            1)
                read -r -p "请输入内存大小 (GiB) [${mem_g}]: " new_mem
                [[ -n "$new_mem" ]] && mem_g="$new_mem"
                ;;
            2)
                read -r -p "请输入网络带宽 (Mbps) [${bw_mbps}]: " new_bw
                [[ -n "$new_bw" ]] && bw_mbps="$new_bw"
                ;;
            3)
                read -r -p "请输入网络延迟 (ms) [${rtt_ms}]: " new_rtt
                [[ -n "$new_rtt" ]] && rtt_ms="$new_rtt"
                ;;
            q|Q)
                info "用户取消操作"
                exit 0
                ;;
            *)
                warning "无效输入，请重试"
                sleep 1
                ;;
        esac
    done
    
    # 执行优化
    local result
    result=$(calculate_buffers "$mem_g" "$bw_mbps" "$rtt_ms")
    read -r bdp_bytes buffer_bytes buffer_mb <<< "$result"
    
    # 处理配置冲突
    handle_conflicts "$backup_dir"
    
    # 加载BBR模块
    if command -v modprobe >/dev/null 2>&1; then
        modprobe tcp_bbr 2>/dev/null || warning "BBR模块加载失败，可能已内置"
    fi
    
    # 应用优化配置
    apply_optimization "$buffer_bytes" "$buffer_mb" "$mem_g" "$bw_mbps" "$rtt_ms" "$bdp_bytes"
    
    # 显示结果
    show_results "$mem_g" "$bw_mbps" "$rtt_ms" "$bdp_bytes" "$buffer_mb"
    
    success "TCP优化配置完成！"
    info "重启系统后所有优化将完全生效"
}

# 执行主函数
main "$@"