#!/bin/bash

# Auto N2 Handover Test Script
# 自动启动两个基站和一个UE，完成N2切换测试

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

# 获取脚本目录
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
UERANSIM_DIR="$(dirname "$SCRIPT_DIR")"

# 配置文件路径
GNB1_CONFIG="$UERANSIM_DIR/config/free5gc-gnb.yaml"
GNB2_CONFIG="$UERANSIM_DIR/config/free5gc-gnb2.yaml"
UE_CONFIG="$UERANSIM_DIR/config/free5gc-ue.yaml"

# 日志目录
LOG_DIR="$UERANSIM_DIR/logs"
mkdir -p "$LOG_DIR"

# PID文件
GNB1_PID_FILE="$LOG_DIR/gnb1.pid"
GNB2_PID_FILE="$LOG_DIR/gnb2.pid"
UE_PID_FILE="$LOG_DIR/ue.pid"

# 清理函数
cleanup() {
    log_step "开始清理进程..."
    
    # 停止UE
    if [ -f "$UE_PID_FILE" ]; then
        UE_PID=$(cat "$UE_PID_FILE")
        if kill -0 "$UE_PID" 2>/dev/null; then
            log_info "停止UE进程 (PID: $UE_PID)"
            sudo kill "$UE_PID" || true
        fi
        rm -f "$UE_PID_FILE"
    fi
    
    # 停止gNB进程
    if [ -f "$GNB1_PID_FILE" ]; then
        GNB1_PID=$(cat "$GNB1_PID_FILE")
        if kill -0 "$GNB1_PID" 2>/dev/null; then
            log_info "停止gNB1进程 (PID: $GNB1_PID)"
            sudo kill "$GNB1_PID" || true
        fi
        rm -f "$GNB1_PID_FILE"
    fi
    
    if [ -f "$GNB2_PID_FILE" ]; then
        GNB2_PID=$(cat "$GNB2_PID_FILE")
        if kill -0 "$GNB2_PID" 2>/dev/null; then
            log_info "停止gNB2进程 (PID: $GNB2_PID)"
            sudo kill "$GNB2_PID" || true
        fi
        rm -f "$GNB2_PID_FILE"
    fi
    
    # 清理所有相关进程
    sudo pkill -f "nr-gnb" || true
    sudo pkill -f "nr-ue" || true
    
    log_info "清理完成"
}

# 检查进程状态
check_process() {
    local pid=$1
    local name=$2
    if kill -0 "$pid" 2>/dev/null; then
        log_info "$name 运行正常 (PID: $pid)"
        return 0
    else
        log_error "$name 进程异常退出"
        return 1
    fi
}

# 等待gNB连接到AMF
wait_gnb_ready() {
    local gnb_name=$1
    local log_file=$2
    local max_wait=30
    local count=0
    
    log_info "等待 $gnb_name 连接到AMF..."
    
    while [ $count -lt $max_wait ]; do
        if grep -q "NG Setup procedure is successful" "$log_file" 2>/dev/null; then
            log_info "$gnb_name 成功连接到AMF"
            return 0
        fi
        sleep 1
        count=$((count + 1))
    done
    
    log_error "$gnb_name 连接AMF超时"
    return 1
}

# 等待UE注册完成
wait_ue_registered() {
    local log_file=$1
    local max_wait=30
    local count=0
    
    log_info "等待UE完成注册..."
    
    while [ $count -lt $max_wait ]; do
        if grep -q "Initial Registration is successful" "$log_file" 2>/dev/null; then
            log_info "UE注册成功"
            return 0
        fi
        sleep 1
        count=$((count + 1))
    done
    
    log_error "UE注册超时"
    return 1
}

# 获取UE ID
get_ue_id() {
    local gnb_name=$1
    local ue_list_output
    ue_list_output=$("$UERANSIM_DIR/build/nr-cli" "$gnb_name" -e "ue-list" 2>/dev/null || echo "")
    
    if [ -n "$ue_list_output" ]; then
        echo "$ue_list_output" | grep "ue-id:" | head -1 | sed 's/.*ue-id: *//' | awk '{print $1}'
    else
        echo ""
    fi
}

# 检查handover是否成功
check_handover_success() {
    local source_gnb=$1
    local target_gnb=$2
    local original_ue_id=$3
    
    log_info "检查handover结果..."
    
    # 检查目标gNB是否有UE
    local target_ue_id
    target_ue_id=$(get_ue_id "$target_gnb")
    
    # 检查源gNB的UE数量
    local source_ue_count
    source_ue_count=$("$UERANSIM_DIR/build/nr-cli" "$source_gnb" -e "ue-list" 2>/dev/null | grep -c "ue-id:" || echo "0")
    
    log_info "源gNB ($source_gnb) UE数量: $source_ue_count"
    log_info "目标gNB ($target_gnb) UE ID: ${target_ue_id:-无}"
    
    # 如果目标gNB有UE，认为handover成功（源gNB可能需要更多时间清理）
    if [ -n "$target_ue_id" ]; then
        log_info "✅ Handover成功！"
        log_info "   目标gNB ($target_gnb): UE ID = $target_ue_id"
        
        if [ "$source_ue_count" -eq 0 ]; then
            log_info "   源gNB ($source_gnb): UE已完全清理"
        else
            log_warn "   源gNB ($source_gnb): UE还未完全清理 (数量: $source_ue_count) - 这是正常的，清理可能需要更多时间"
        fi
        return 0
    else
        log_error "❌ Handover失败！"
        log_error "   源gNB UE数量: $source_ue_count"
        log_error "   目标gNB UE ID: ${target_ue_id:-无}"
        return 1
    fi
}

# 主函数
main() {
    log_step "=== 自动N2 Handover测试开始 ==="
    
    # 设置退出时清理
    trap cleanup EXIT
    
    # 1. 清理现有进程
    log_step "步骤1: 清理现有进程"
    cleanup
    sleep 2
    
    # 2. 启动gNB1 (目标gNB)
    log_step "步骤2: 启动gNB1 (192.168.8.117)"
    sudo "$UERANSIM_DIR/build/nr-gnb" -c "$GNB1_CONFIG" > "$LOG_DIR/gnb1.log" 2>&1 &
    GNB1_PID=$!
    echo "$GNB1_PID" > "$GNB1_PID_FILE"
    log_info "gNB1已启动 (PID: $GNB1_PID)"
    
    # 等待gNB1连接AMF
    if ! wait_gnb_ready "gNB1" "$LOG_DIR/gnb1.log"; then
        log_error "gNB1启动失败"
        exit 1
    fi
    
    # 3. 启动gNB2 (源gNB)
    log_step "步骤3: 启动gNB2 (192.168.8.118)"
    sleep 3
    sudo "$UERANSIM_DIR/build/nr-gnb" -c "$GNB2_CONFIG" > "$LOG_DIR/gnb2.log" 2>&1 &
    GNB2_PID=$!
    echo "$GNB2_PID" > "$GNB2_PID_FILE"
    log_info "gNB2已启动 (PID: $GNB2_PID)"
    
    # 等待gNB2连接AMF
    if ! wait_gnb_ready "gNB2" "$LOG_DIR/gnb2.log"; then
        log_error "gNB2启动失败"
        exit 1
    fi
    
    # 4. 启动UE并连接到gNB2
    log_step "步骤4: 启动UE并连接到gNB2"
    sleep 3
    sudo "$UERANSIM_DIR/build/nr-ue" -c "$UE_CONFIG" > "$LOG_DIR/ue.log" 2>&1 &
    UE_PID=$!
    echo "$UE_PID" > "$UE_PID_FILE"
    log_info "UE已启动 (PID: $UE_PID)"
    
    # 等待UE注册完成
    if ! wait_ue_registered "$LOG_DIR/ue.log"; then
        log_error "UE注册失败"
        exit 1
    fi
    
    # 5. 检查初始状态
    log_step "步骤5: 检查初始UE分布"
    sleep 2
    
    # 确定UE连接到哪个gNB
    log_info "检查gNB1的UE列表..."
    UE_LIST_GNB1=$("$UERANSIM_DIR/build/nr-cli" "UERANSIM-gnb-460-0-1" -e "ue-list" 2>/dev/null || echo "")
    log_info "gNB1 UE列表输出: '$UE_LIST_GNB1'"
    
    log_info "检查gNB2的UE列表..."
    UE_LIST_GNB2=$("$UERANSIM_DIR/build/nr-cli" "UERANSIM-gnb-460-0-16" -e "ue-list" 2>/dev/null || echo "")
    log_info "gNB2 UE列表输出: '$UE_LIST_GNB2'"
    
    UE_ID_GNB1=$(echo "$UE_LIST_GNB1" | grep "ue-id:" | head -1 | sed 's/.*ue-id: *//' | awk '{print $1}')
    UE_ID_GNB2=$(echo "$UE_LIST_GNB2" | grep "ue-id:" | head -1 | sed 's/.*ue-id: *//' | awk '{print $1}')
    
    if [ -n "$UE_ID_GNB1" ]; then
        SOURCE_GNB="UERANSIM-gnb-460-0-1"
        TARGET_GNB="UERANSIM-gnb-460-0-16"
        TARGET_CELL="16"
        UE_ID="$UE_ID_GNB1"
        log_info "UE连接到gNB1，UE ID: $UE_ID"
    elif [ -n "$UE_ID_GNB2" ]; then
        SOURCE_GNB="UERANSIM-gnb-460-0-16"
        TARGET_GNB="UERANSIM-gnb-460-0-1"
        TARGET_CELL="1"
        UE_ID="$UE_ID_GNB2"
        log_info "UE连接到gNB2，UE ID: $UE_ID"
    else
        log_error "未找到UE连接"
        exit 1
    fi
    
    # 6. 执行handover
    log_step "步骤6: 执行N2 Handover"
    log_info "从 $SOURCE_GNB 切换UE $UE_ID 到目标Cell $TARGET_CELL"
    
    HANDOVER_OUTPUT=$("$UERANSIM_DIR/build/nr-cli" "$SOURCE_GNB" -e "handover $UE_ID $TARGET_CELL" 2>&1)
    log_info "Handover命令输出: $HANDOVER_OUTPUT"
    
    # 7. 等待handover完成
    log_step "步骤7: 等待handover完成"
    sleep 15  # 等待handover完成 (8秒清理时间 + 7秒缓冲)
    
    # 8. 验证handover结果
    log_step "步骤8: 验证handover结果"
    if check_handover_success "$SOURCE_GNB" "$TARGET_GNB" "$UE_ID"; then
        log_step "=== ✅ N2 Handover测试成功完成 ==="
        
        # 显示最终状态
        echo ""
        log_info "最终UE分布:"
        echo "gNB1 (UERANSIM-gnb-460-0-1):"
        "$UERANSIM_DIR/build/nr-cli" "UERANSIM-gnb-460-0-1" -e "ue-list" 2>/dev/null || echo "  无UE连接"
        echo ""
        echo "gNB2 (UERANSIM-gnb-460-0-16):"
        "$UERANSIM_DIR/build/nr-cli" "UERANSIM-gnb-460-0-16" -e "ue-list" 2>/dev/null || echo "  无UE连接"
        
        # 保持运行以便观察
        echo ""
        log_info "测试完成！按Ctrl+C退出并清理进程"
        while true; do
            sleep 5
            # 检查进程状态
            if ! check_process "$GNB1_PID" "gNB1" || ! check_process "$GNB2_PID" "gNB2" || ! check_process "$UE_PID" "UE"; then
                log_warn "发现进程异常，退出测试"
                break
            fi
        done
    else
        log_step "=== ❌ N2 Handover测试失败 ==="
        exit 1
    fi
}

# 检查是否以root权限运行
if [ "$EUID" -ne 0 ]; then
    echo -e "${RED}[ERROR]${NC} 请使用sudo运行此脚本"
    exit 1
fi

# 检查构建目录
if [ ! -f "$UERANSIM_DIR/build/nr-gnb" ] || [ ! -f "$UERANSIM_DIR/build/nr-ue" ]; then
    log_error "找不到UERANSIM构建文件，请先运行make编译"
    exit 1
fi

# 检查配置文件
for config in "$GNB1_CONFIG" "$GNB2_CONFIG" "$UE_CONFIG"; do
    if [ ! -f "$config" ]; then
        log_error "找不到配置文件: $config"
        exit 1
    fi
done

# 运行主函数
main "$@"