#!/bin/bash

# 默认使用 Ubuntu 镜像，执行脚本时可传参指定，例如: ./test-ai-container-net.sh alpine
IMAGE=${1:-"ubuntu:22.04"}

echo "🚀 正在拉起临时容器 ($IMAGE) 进行网络连通性测试..."

docker run --rm -i "$IMAGE" sh << 'EOF'

# --- 以下代码全都在容器内部执行 ---
export DEBIAN_FRONTEND=noninteractive
OS_NAME=$(cat /etc/os-release 2>/dev/null | grep PRETTY_NAME | cut -d'"' -f2 || echo "Unknown OS")

echo "================================================================"
echo " 🌐 容器内 AI 网络与环境测试 (系统: $OS_NAME)"
echo " ⚙️ 注入代理配置: HTTP_PROXY=${HTTP_PROXY:-未设置}, HTTPS_PROXY=${HTTPS_PROXY:-未设置}"
echo "================================================================"

# ----------------------------------------------------------------
# 阶段 1: 基础包管理器 (apt/apk) 连通性测试
# ----------------------------------------------------------------
echo -e "\n[阶段 1] 测试系统包管理器能否正常拉取软件源..."

PKG_MGR_STATUS="FAIL"

if command -v apt-get >/dev/null 2>&1; then
    echo -n "⏳ 正在执行 apt-get update... "
    # 捕获错误输出到临时文件，以便失败时展示给用户
    if apt-get update -qq >/tmp/apt_error.log 2>&1; then
        echo -e "\033[0;32m[OK]\033[0m"
        echo "📦 正在静默安装 curl 和 awk 以进行后续测试..."
        apt-get install -y curl awk >/dev/null 2>&1
        PKG_MGR_STATUS="OK"
    else
        echo -e "\033[0;31m[FAIL]\033[0m"
        echo -e "\033[0;33m⚠️ 错误详情:\033[0m"
        cat /tmp/apt_error.log
    fi
elif command -v apk >/dev/null 2>&1; then
    echo -n "⏳ 正在执行 apk update... "
    if apk update >/tmp/apk_error.log 2>&1; then
        echo -e "\033[0;32m[OK]\033[0m"
        echo "📦 正在静默安装 curl 和 awk 以进行后续测试..."
        apk add --no-cache curl awk >/dev/null 2>&1
        PKG_MGR_STATUS="OK"
    else
        echo -e "\033[0;31m[FAIL]\033[0m"
        echo -e "\033[0;33m⚠️ 错误详情:\033[0m"
        cat /tmp/apk_error.log
    fi
else
    echo -e "\033[0;33m[SKIP] 未检测到 apt 或 apk，跳过包管理器测试。\033[0m"
fi

# 如果包管理器更新失败，直接阻断退出
if [ "$PKG_MGR_STATUS" = "FAIL" ]; then
    echo "----------------------------------------------------------------"
    echo -e "\033[0;31m❌ 致命错误: 容器内无法连接系统官方软件源。\033[0m"
    echo "👉 排查建议:"
    echo "   1. 检查宿主机 Docker 是否被配置了错误的全局代理/镜像加速器。"
    echo "   2. 检查宿主机的防火墙(iptables)是否拦截了 Docker 网桥的 53 端口(DNS)或 80/443 端口。"
    exit 1
fi

if ! command -v curl >/dev/null 2>&1; then
    echo -e "\033[0;31m❌ 致命错误: curl 安装失败，无法继续测试。\033[0m"
    exit 1
fi

# ----------------------------------------------------------------
# 阶段 2: 目标网站连通性与延迟测试
# ----------------------------------------------------------------
echo -e "\n[阶段 2] 测试 AI 开发常用目标节点连通性..."
printf "%-32s %-15s %-15s\n" "测试目标" "状态" "连接耗时"
echo "----------------------------------------------------------------"

TARGETS="
HuggingFace_Web|https://huggingface.co
HuggingFace_API|https://huggingface.co/api/models
GitHub_Web|https://github.com
GitHub_Raw|https://raw.githubusercontent.com
OpenAI_API|https://api.openai.com/v1/models
Anthropic_API|https://api.anthropic.com/v1/messages
Google_Gemini|https://generativelanguage.googleapis.com
Docker_Hub|https://registry-1.docker.io/v2/
PyPI_Python源|https://pypi.org
"

for item in $TARGETS; do
    [ -z "$item" ] && continue 

    name=$(echo "$item" | cut -d'|' -f1)
    url=$(echo "$item" | cut -d'|' -f2)

    result=$(curl -s -o /dev/null -w "%{http_code}|%{time_connect}" -m 5 "$url")
    
    http_code=$(echo "$result" | cut -d'|' -f1)
    time_connect=$(echo "$result" | cut -d'|' -f2)

    if [ "$time_connect" != "0.000000" ]; then
        time_ms=$(awk "BEGIN {printf \"%.0f\", $time_connect * 1000}")" ms"
    else
        time_ms="超时"
    fi

    if [ "$http_code" = "000" ]; then
        printf "%-32s \033[0;31m%-15s\033[0m %-15s\n" "$name" "FAIL(超时/断网)" "$time_ms"
    elif [ "$http_code" = "401" ] || [ "$http_code" = "403" ] || [ "$http_code" = "404" ]; then
        printf "%-32s \033[0;32m%-15s\033[0m %-15s\n" "$name" "OK($http_code)" "$time_ms"
    else
        printf "%-32s \033[0;32m%-15s\033[0m %-15s\n" "$name" "OK($http_code)" "$time_ms"
    fi
done

echo "----------------------------------------------------------------"
EOF