#!/bin/bash

# ==========================================
# 利群主机 LeiKwan Host - WireGuard Auto Setup
# ==========================================

# --- 颜色定义 / Colors ---
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
WHITE='\033[1;37m'
NC='\033[0m' # No Color

# --- 全局变量 (默认值) / Global Variables (Default) ---
# 这些变量会在 ask_interface 函数中根据用户输入更新
WG_IFACE="wg0"
WG_CONF="/etc/wireguard/${WG_IFACE}.conf"
PRIVATE_KEY_FILE="/etc/wireguard/${WG_IFACE}_privatekey"
PUBLIC_KEY_FILE="/etc/wireguard/${WG_IFACE}_publickey"
SERVICE_NAME="wg-quick@${WG_IFACE}"

# --- 辅助函数 / Helper Functions ---

show_banner() {
    clear
    echo -e "${CYAN}"
    echo "  _          _  __                        _   _           _   "
    echo " | |    ___ (_) |/ /      ____ _ _ __    | | | | ___  ___| |_ "
    echo " | |   / _ \| | ' / \ /\ / / _\` | '_ \   | |_| |/ _ \/ __| __|"
    echo " | |__|  __/| | . \  V  V / (_| | | | |  |  _  | (_) \__ \ |_ "
    echo " |_____\___||_|_|\_\ \_/\_/ \__,_|_| |_|  |_| |_|\___/|___/\__|"
    echo -e "${WHITE}           利 群 主 機  -  L e i K w a n   H o s t${NC}"
    echo -e "${CYAN} ==========================================================${NC}"
    echo -e "${WHITE}           WireGuard Auto Setup Script v2.2${NC}"
    echo -e "${CYAN} ==========================================================${NC}"
    echo ""
}

print_info() { echo -e "${CYAN}[INFO]${NC} $1"; }
print_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
print_warn() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
print_error() { echo -e "${RED}[ERROR]${NC} $1"; }

print_card() {
    local title=$1
    shift
    echo -e "\n${WHITE}┌────────────────────────────────────────────────────────┐${NC}"
    echo -e "${WHITE}│ ${YELLOW}$title ${NC}"
    echo -e "${WHITE}├────────────────────────────────────────────────────────┤${NC}"
    for line in "$@"; do
        printf "${WHITE}│${NC} %-54s ${WHITE}│${NC}\n" "$line"
    done
    echo -e "${WHITE}└────────────────────────────────────────────────────────┘${NC}\n"
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        print_error "Please run as root: sudo ./auto-wg.sh"
        exit 1
    fi
}

# --- 设置函数 / Setup Functions ---

# 新增：自定义接口名称逻辑
ask_interface() {
    echo -e "${YELLOW}? 自定義接口名稱 / Custom Interface Name${NC}"
    echo -e "   默認 / Default: ${GREEN}wg0${NC}"
    read -p "   Enter Name (e.g. wg0, wg1, tun0): " input_iface
    
    if [[ -n "$input_iface" ]]; then
        WG_IFACE="$input_iface"
    fi

    # 更新所有相关路径和服务名
    WG_CONF="/etc/wireguard/${WG_IFACE}.conf"
    PRIVATE_KEY_FILE="/etc/wireguard/${WG_IFACE}_privatekey"
    PUBLIC_KEY_FILE="/etc/wireguard/${WG_IFACE}_publickey"
    SERVICE_NAME="wg-quick@${WG_IFACE}"
    
    print_info "Selected Interface: ${GREEN}${WG_IFACE}${NC}"
    print_info "Config Path: ${WG_CONF}"
    echo ""
}

install_wg() {
    if ! command -v wg &> /dev/null; then
        print_info "Installing WireGuard..."
        if [ -x "$(command -v apt)" ]; then
            apt update -qq && apt install -y wireguard curl > /dev/null
        elif [ -x "$(command -v yum)" ]; then
            yum install -y epel-release > /dev/null
            yum install -y wireguard-tools curl > /dev/null
        else
            print_error "OS not supported. Please install wireguard & curl manually."
            exit 1
        fi
        print_success "WireGuard installed."
    else
        print_info "WireGuard is already installed."
    fi
}

gen_keys() {
    umask 077
    mkdir -p /etc/wireguard
    if [ ! -f "$PRIVATE_KEY_FILE" ]; then
        wg genkey | tee "$PRIVATE_KEY_FILE" | wg pubkey > "$PUBLIC_KEY_FILE"
        print_success "Keys generated."
    else
        print_info "Existing keys found, using them."
    fi
    PRIV_KEY=$(cat "$PRIVATE_KEY_FILE")
    PUB_KEY=$(cat "$PUBLIC_KEY_FILE")
}

# --- 核心流程 / Core Logic ---

logic_A() {
    echo -e "${WHITE}>>> Mode: ${CYAN}A. 利群主機 (內網/發起端) / Intranet Host${NC}"
    echo ""
    
    echo -e "${YELLOW}? 問 / Ask:${NC} B機器 (雲端) 是否已部署完畢？ / Is Server B ready?"
    read -p "   (y/n): " b_ready

    install_wg
    gen_keys

    echo ""
    read -p "-> 本機內網IP / Local Intranet IP [Default 10.198.1.1]: " local_ip
    local_ip=${local_ip:-10.198.1.1}

    # 生成基础配置
    cat > $WG_CONF <<EOF
[Interface]
# Generate by LeiKwan Host Script
PrivateKey = $PRIV_KEY
Address = $local_ip/24
MTU = 1280
EOF

    if [[ "$b_ready" == "y" || "$b_ready" == "Y" ]]; then
        echo -e "\n${GREEN}>>> Importing Config from B...${NC}"
        read -p "   B's IP:Port : " b_endpoint
        read -p "   B's PUBLIC KEY : " b_pubkey
        
        # 自动推算
        base_ip=$(echo $local_ip | cut -d'.' -f1-3)
        last_octet=$(echo $local_ip | cut -d'.' -f4)
        target_octet=$((last_octet + 1))
        auto_b_ip="${base_ip}.${target_octet}"

        read -p "   B's Internal IP [Default $auto_b_ip]: " b_internal_ip
        b_internal_ip=${b_internal_ip:-$auto_b_ip}

        cat >> $WG_CONF <<EOF

[Peer]
PublicKey = $b_pubkey
Endpoint = $b_endpoint
AllowedIPs = $b_internal_ip/32
PersistentKeepalive = 25
EOF
        systemctl enable $SERVICE_NAME >/dev/null 2>&1
        systemctl restart $SERVICE_NAME
        print_success "Done! Trying to ping B..."
        ping -c 3 -W 1 "$b_internal_ip"

    else
        print_success "Basic config generated."
        print_card "Copy to Machine B / 複製到 B 機器" \
            "IP_ADDR   : $local_ip" \
            "PUBLICKEY : $PUB_KEY"
            
        print_warn "Next Step / 下一步:"
        echo "1. Go to Machine B -> Run Script -> Choose 'B'."
        echo "2. Come back here (A) -> Run Script -> Choose 'C'."
    fi
}

logic_B() {
    echo -e "${WHITE}>>> Mode: ${CYAN}B. 其他機器 (雲端/公網) / Cloud Server${NC}"
    echo ""

    echo -e "${YELLOW}? 問 / Ask:${NC} A機器 (利群主機) 是否已完成第一步？ / Is Host A ready?"
    read -p "   (y/n): " a_done

    if [[ "$a_done" != "y" && "$a_done" != "Y" ]]; then
        print_error "Please setup Host A first."
        exit 0
    fi

    echo -e "\n${GREEN}>>> Input Info from A${NC}"
    read -p "   Paste A's PUBLIC KEY: " a_pubkey
    read -p "   Paste A's Internal IP (e.g. 10.198.1.1): " a_internal_ip

    if [[ -z "$a_pubkey" || -z "$a_internal_ip" ]]; then
        print_error "Empty input. Exit."
        exit 1
    fi

    install_wg
    gen_keys

    read -p "   Listen Port [Default $(shuf -i 8000-9000 -n 1)]: " custom_port
    if [[ -n "$custom_port" && "$custom_port" =~ ^[0-9]{1,5}$ ]] && [ "$custom_port" -ge 1 ] && [ "$custom_port" -le 65535 ]; then
        listen_port="$custom_port"
    else
        listen_port=$(shuf -i 8000-9000 -n 1)
        print_info "Using random port: $listen_port"
    fi
    
    base_ip=$(echo $a_internal_ip | cut -d'.' -f1-3)
    last_octet=$(echo $a_internal_ip | cut -d'.' -f4)
    my_octet=$((last_octet + 1))
    default_my_ip="${base_ip}.${my_octet}"

    read -p "   Local Intranet IP [Default $default_my_ip]: " my_ip
    my_ip=${my_ip:-$default_my_ip}

    print_info "Detecting Public IP (IPv4)..."
    # 修改点：增加 -4 参数强制获取 IPv4
    public_ip=$(curl -4 -s --max-time 3 ifconfig.me)
    
    if [[ -z "$public_ip" ]]; then
        print_warn "IPv4 detection failed, trying default..."
        public_ip=$(curl -s --max-time 3 ifconfig.me)
    fi

    if [[ -z "$public_ip" ]]; then
        read -p "   Failed to detect IP. Please input Public IP manually: " public_ip
    else
        print_info "Detected IP: $public_ip"
    fi

    cat > $WG_CONF <<EOF
[Interface]
# Generate by LeiKwan Host Script
PrivateKey = $PRIV_KEY
Address = $my_ip/24
ListenPort = $listen_port
MTU = 1280
PreUp = sysctl -w net.ipv4.ip_forward=1

[Peer]
PublicKey = $a_pubkey
AllowedIPs = $a_internal_ip/32
EOF

    systemctl enable $SERVICE_NAME >/dev/null 2>&1
    systemctl restart $SERVICE_NAME

    print_success "Server B Configured."
    print_card "Copy back to Machine A (Step C) / 複製回 A 機器" \
        "IP:Port   : ${public_ip}:${listen_port}" \
        "PUBLICKEY : ${PUB_KEY}" \
        "Address   : ${my_ip}"
        
    print_warn "Now go back to Machine A and choose 'C'. / 請回到 A 機器選 C"
}

logic_C() {
    echo -e "${WHITE}>>> Mode: ${CYAN}C. 最終握手 / Final Handshake (On A)${NC}"
    
    if [ ! -f "$WG_CONF" ]; then
        print_error "Config file ($WG_CONF) not found. Run Step A first."
        exit 1
    fi

    echo -e "\n${GREEN}>>> Input Info from B${NC}"
    read -p "   Paste B's IP:Port : " b_endpoint
    read -p "   Paste B's PUBLIC KEY : " b_pubkey
    read -p "   Paste B's Address (Intranet IP) [Default 10.198.1.2]: " b_ip
    b_ip=${b_ip:-10.198.1.2}

    if grep -q "$b_pubkey" "$WG_CONF"; then
        print_warn "Peer already exists, appending anyway..."
    fi

    cat >> $WG_CONF <<EOF

[Peer]
PublicKey = $b_pubkey
Endpoint = $b_endpoint
AllowedIPs = $b_ip/32
PersistentKeepalive = 25
EOF

    print_info "Restarting WireGuard ($WG_IFACE)..."
    systemctl restart $SERVICE_NAME
    
    echo -e "\n${GREEN}==========================================${NC}"
    echo -e "${GREEN}   利群主機 WireGuard 隧道已建立 ($WG_IFACE)   ${NC}"
    echo -e "${GREEN}==========================================${NC}"
    echo -e "Testing connection (Ping $b_ip)..."
    ping -c 3 -W 1 "$b_ip"
}

# --- 主程序入口 / Main Entry ---
check_root
show_banner
ask_interface  # <--- 在此处优先询问接口名称

echo -e "Select Option / 選擇模式:"
echo -e "${CYAN}1.${NC} A. 利群主機 (內網) / Intranet Host"
echo -e "${CYAN}2.${NC} B. 其他機器 (雲端) / Cloud Server"
echo -e "${CYAN}3.${NC} C. 完成 A/B 後的最後一步 / Final Step (On A)"
echo ""
read -p "Input [1/2/3] or [A/B/C]: " choice

case "$choice" in
    1|a|A) logic_A ;;
    2|b|B) logic_B ;;
    3|c|C) logic_C ;;
    *) print_error "Invalid option." ;;
esac