#!/bin/bash
# ============================================================================
# srg-diagnose — 远程环境诊断
# ============================================================================
# 逐层检查远程代理环境，输出诊断信息和修复建议。
#
# 参数:
#   $1 — 代理端口
# ============================================================================

set -uo pipefail

PORT="${1:?缺少参数: 代理端口}"
ADDR="127.0.0.1:$PORT"

# Supported IDE server directories (add more to support other IDEs, e.g. .vscode-server)
IDE_SERVER_DIRS=(".antigravity-ide-server" ".antigravity-server")

PASS=0
FAIL=0
WARN=0

pass()  { echo -e "  \033[0;32m✓\033[0m $1"; PASS=$((PASS + 1)); }
fail()  { echo -e "  \033[0;31m✗\033[0m $1"; FAIL=$((FAIL + 1)); }
warn_() { echo -e "  \033[1;33m⚠\033[0m $1"; WARN=$((WARN + 1)); }
hint()  { echo -e "    \033[2m→ $1\033[0m"; }

echo "SRG 远程环境诊断"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# ============================================================================
# 1. SSH 隧道
# ============================================================================
echo ""
echo "  [1] SSH 反向隧道 (:$PORT)"

if timeout 3 bash -c "echo > /dev/tcp/127.0.0.1/$PORT" 2>/dev/null; then
    pass "端口 $PORT 可达"
else
    fail "端口 $PORT 不可达"
    hint "检查本地代理是否在 $PORT 端口运行"
    hint "检查 SSH 连接是否使用了 RemoteForward"
    hint "运行: ssh -v <host> 查看是否有 'remote forward success'"

    # 检查是否有其它进程占用这个端口
    PID=$(lsof -ti tcp:$PORT 2>/dev/null || ss -tlnp "sport = :$PORT" 2>/dev/null | tail -1 | grep -oP 'pid=\K\d+' || true)
    if [ -n "$PID" ]; then
        PROC=$(ps -p "$PID" -o comm= 2>/dev/null || echo "未知")
        warn_ "端口 $PORT 被进程 $PROC (PID:$PID) 占用，可能导致 RemoteForward 失败"
        hint "运行: kill $PID 或更换端口"
    fi
fi

# ============================================================================
# 2. 代理协议测试（仅在隧道可达时）
# ============================================================================
echo ""
echo "  [2] 代理协议"

if timeout 2 bash -c "echo > /dev/tcp/127.0.0.1/$PORT" 2>/dev/null; then
    # HTTP
    CODE=$(curl -x http://$ADDR https://www.google.com -s -o /dev/null -w "%{http_code}" --connect-timeout 10 2>/dev/null || echo "000")
    if [[ "$CODE" =~ ^(200|301|302)$ ]]; then
        pass "HTTP 代理 → Google (HTTP $CODE)"
    else
        fail "HTTP 代理 → Google (HTTP $CODE)"
        if [ "$CODE" = "000" ]; then
            hint "连接超时，本地代理可能不支持 HTTP CONNECT"
        else
            hint "代理返回异常状态码"
        fi
    fi

    # SOCKS5
    CODE=$(curl -x socks5://$ADDR https://www.google.com -s -o /dev/null -w "%{http_code}" --connect-timeout 10 2>/dev/null || echo "000")
    if [[ "$CODE" =~ ^(200|301|302)$ ]]; then
        pass "SOCKS5 代理 → Google (HTTP $CODE)"
    else
        fail "SOCKS5 代理 → Google (HTTP $CODE)"
        if [ "$CODE" = "000" ]; then
            hint "本地代理可能不支持 SOCKS5 协议"
        fi
    fi

    # DNS 解析测试
    RESOLVE=$(curl -x http://$ADDR https://httpbin.org/ip -s --connect-timeout 10 2>/dev/null || echo "")
    if echo "$RESOLVE" | grep -q "origin"; then
        pass "DNS 通过代理解析正常"
    else
        warn_ "DNS 解析测试未通过 (httpbin.org 不可达)"
    fi
else
    warn_ "跳过代理协议测试 (隧道不可达)"
fi

# ============================================================================
# 3. SRG 工具完整性
# ============================================================================
echo ""
echo "  [3] SRG 工具集"

ALL_TOOLS_OK=true
for tool in srg-on srg-off srg-shell srg-proxy srg-status; do
    if [ -f "$HOME/bin/$tool" ]; then
        if [ -x "$HOME/bin/$tool" ]; then
            pass "$tool"
        else
            fail "$tool (存在但不可执行)"
            hint "运行: chmod +x ~/bin/$tool"
            ALL_TOOLS_OK=false
        fi
    else
        fail "$tool (缺失)"
        ALL_TOOLS_OK=false
    fi
done

if ! $ALL_TOOLS_OK; then
    hint "重新运行 srg setup 修复"
fi

# ============================================================================
# 4. mgraftcp-fakedns
# ============================================================================
echo ""
echo "  [4] mgraftcp-fakedns"

MGRAFTCP="$HOME/bin/mgraftcp-fakedns"
if [ -x "$MGRAFTCP" ]; then
    pass "mgraftcp-fakedns 已安装 ($MGRAFTCP)"

    # 检查动态链接
    if command -v ldd &>/dev/null; then
        MISSING_LIBS=$(ldd "$MGRAFTCP" 2>/dev/null | grep "not found" || true)
        if [ -n "$MISSING_LIBS" ]; then
            fail "缺少动态库:"
            echo "$MISSING_LIBS" | while read -r line; do
                hint "$line"
            done
        else
            pass "动态库依赖完整"
        fi
    fi

    # 检查 libdnsredir
    LIB=$(find "$HOME/bin" -name "libdnsredir*.so" -type f 2>/dev/null | head -1 || true)
    if [ -n "$LIB" ]; then
        pass "libdnsredir.so 已安装 ($(basename "$LIB"))"
    else
        warn_ "libdnsredir.so 未找到 (DNS 重定向可能不工作)"
        hint "重新运行 srg setup 或手动编译 graftcp"
    fi
else
    fail "mgraftcp-fakedns 未安装"
    hint "重新运行 srg setup，或手动编译:"
    hint "  git clone https://github.com/hmgle/graftcp && cd graftcp && make"
    hint "  cp local/mgraftcp ~/bin/mgraftcp-fakedns && chmod +x ~/bin/mgraftcp-fakedns"
fi

# ============================================================================
# 5. Shell 环境
# ============================================================================
echo ""
echo "  [5] Shell 环境"

# PATH 检查
if echo "$PATH" | tr ':' '\n' | grep -q "$HOME/bin"; then
    pass "~/bin 在 PATH 中"
else
    fail "~/bin 不在 PATH 中"
    hint "运行: export PATH=\"\$HOME/bin:\$PATH\""
    hint "或重新 source ~/.bashrc"
fi

# bashrc 函数检查
if grep -q "srg-on()" "$HOME/.bashrc" 2>/dev/null; then
    pass ".bashrc 已注入 SRG 函数"
else
    warn_ ".bashrc 未注入 SRG 函数 (srg-on/off 需手动调用完整路径)"
fi

# 当前 shell 代理状态
if [ -n "${HTTP_PROXY:-}" ]; then
    pass "当前 shell 代理: 已开启 ($HTTP_PROXY)"
else
    pass "当前 shell 代理: 未开启 (正常，按需用 srg-on 开启)"
fi

# ============================================================================
# 6. Language Server Wrapper
# ============================================================================
echo ""
echo "  [6] Language Server Wrapper"

_find_args=()
for _d in "${IDE_SERVER_DIRS[@]}"; do
    if [ -d "$HOME/$_d" ]; then
        _find_args+=("$HOME/$_d")
    fi
done
if [ ${#_find_args[@]} -eq 0 ]; then
    _find_args+=("$HOME/.antigravity-ide-server" "$HOME/.antigravity-server")
fi
LS_TARGETS=$(find "${_find_args[@]}" \
    -type f -name "language_server_linux_*" 2>/dev/null | grep -v "\.bak$" || true)

if [ -z "$LS_TARGETS" ]; then
    warn_ "未找到 Language Server 二进制 (IDE 可能未安装)"
else
    while IFS= read -r LS; do
        [ -z "$LS" ] && continue
        NAME=$(basename "$LS")
        BAK="${LS}.bak"

        if head -1 "$LS" 2>/dev/null | grep -q "^#!/bin/bash"; then
            if [ -f "$BAK" ]; then
                pass "$NAME → wrapper 已配置 (.bak 备份存在)"
            else
                fail "$NAME → 是脚本但无 .bak 备份 (可能异常)"
                hint "建议重新运行 srg setup"
            fi
        else
            fail "$NAME → 未配置 wrapper (LS 流量不会走代理)"
            hint "运行 srg setup <host> 配置 wrapper"
        fi
    done <<< "$LS_TARGETS"
fi

# ============================================================================
# 7. 系统环境
# ============================================================================
echo ""
echo "  [7] 系统环境"

# 架构
ARCH=$(uname -m)
pass "架构: $ARCH ($(uname -s) $(uname -r))"

# curl 版本
if command -v curl &>/dev/null; then
    CURL_VER=$(curl --version 2>/dev/null | head -1 | awk '{print $2}')
    pass "curl: $CURL_VER"
else
    fail "curl 未安装"
fi

# 磁盘空间
DISK_FREE=$(df -h "$HOME" 2>/dev/null | tail -1 | awk '{print $4}')
if [ -n "$DISK_FREE" ]; then
    pass "磁盘可用空间: $DISK_FREE"
fi

# ============================================================================
# 汇总
# ============================================================================
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
TOTAL=$((PASS + FAIL + WARN))
echo -e "  通过: \033[0;32m$PASS\033[0m / $TOTAL    失败: \033[0;31m$FAIL\033[0m    警告: \033[1;33m$WARN\033[0m"

if [ $FAIL -eq 0 ]; then
    echo -e "  \033[0;32m环境正常\033[0m"
elif [ $FAIL -le 2 ]; then
    echo -e "  \033[1;33m存在少量问题，请参照上方提示修复\033[0m"
else
    echo -e "  \033[0;31m存在较多问题，建议重新运行 srg setup\033[0m"
fi
