#!/bin/bash
# ==============================================================================
# Arthas MCP 服务停止脚本
# 多级优雅退出：HTTP API -> 官方客户端(Telnet) -> Attach 进程退出 -> 端口释放核验
# ==============================================================================
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PACKAGE_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"

# 1. 智能参数识别 (支持 stop-mcp.sh [PORT|PID] [PID|PORT] [PASSWORD])
ARG1="$1"
ARG2="$2"
ARG3="$3"

HTTP_PORT=""
TARGET_PID=""
AUTH_PASSWORD=""

# 辅助函数：通过端口查找监听进程 PID
find_pid_by_port() {
    local port="$1"
    local pid=""
    if command -v lsof &>/dev/null; then
        pid=$(lsof -t -s TCP:LISTEN -i TCP:"${port}" 2>/dev/null | head -n 1)
    fi
    if [[ -z "$pid" ]] && command -v fuser &>/dev/null; then
        pid=$(fuser "${port}/tcp" 2>/dev/null | awk '{print $1}')
    fi
    if [[ -z "$pid" ]] && command -v ss &>/dev/null; then
        pid=$(ss -tlnp "sport = :${port}" 2>/dev/null | grep -oP 'pid=\K[0-9]+' | head -n 1)
    fi
    if [[ -z "$pid" ]] && command -v netstat &>/dev/null; then
        pid=$(netstat -tlpn 2>/dev/null | grep ":${port} " | awk '{print $7}' | cut -d/ -f1 | head -n 1)
    fi
    echo "$pid"
}

# 判断输入是 PID 还是 PORT
if [[ -n "$ARG1" ]] && kill -0 "$ARG1" 2>/dev/null; then
    TARGET_PID="$ARG1"
    HTTP_PORT="${ARG2:-8563}"
    AUTH_PASSWORD="${ARG3:-arthas123}"
elif [[ -n "$ARG2" ]] && kill -0 "$ARG2" 2>/dev/null; then
    HTTP_PORT="${ARG1:-8563}"
    TARGET_PID="$ARG2"
    AUTH_PASSWORD="${ARG3:-arthas123}"
else
    HTTP_PORT="${ARG1:-8563}"
    TARGET_PID="$ARG2"
    AUTH_PASSWORD="${ARG3:-arthas123}"
fi

# 若未指定 PID，自动尝试从端口关联获取
if [[ -z "$TARGET_PID" ]]; then
    DETECTED_PID=$(find_pid_by_port "${HTTP_PORT}")
    if [[ -n "$DETECTED_PID" ]]; then
        TARGET_PID="$DETECTED_PID"
        echo ">> 自动定位到端口 ${HTTP_PORT} 对应的 Java 进程 PID: ${TARGET_PID}"
    fi
fi

# 获取 Telnet 端口（默认 3658）
TELNET_PORT="3658"
if [[ -f "${PACKAGE_ROOT}/arthas-bin/arthas.properties" ]]; then
    PROP_TELNET=$(grep -E "^arthas.telnetPort=" "${PACKAGE_ROOT}/arthas-bin/arthas.properties" | cut -d'=' -f2 | tr -d ' \r\n')
    [[ -n "$PROP_TELNET" ]] && TELNET_PORT="$PROP_TELNET"
fi

echo ">> 正在请求停止 Arthas 诊断会话 (HTTP端口: ${HTTP_PORT}, Telnet端口: ${TELNET_PORT})..."

STOP_SENT=0

# --- 第 1 级: 尝试通过 Arthas HTTP API 发送 stop 命令 ---
if command -v curl &>/dev/null; then
    echo ">> [1/3] 尝试通过 HTTP API (/api) 触发 stop..."
    CURL_AUTH=()
    if [[ -n "${AUTH_PASSWORD}" ]]; then
        CURL_AUTH=(-H "Authorization: Bearer ${AUTH_PASSWORD}")
    fi

    # 发送 stop 命令并限制超时
    HTTP_RES=$(curl --noproxy "*" --connect-timeout 2 -m 4 -s -X POST "http://127.0.0.1:${HTTP_PORT}/api" \
        -H "Content-Type: application/json" \
        "${CURL_AUTH[@]}" \
        -d '{"action":"exec","command":"stop"}' 2>/dev/null || true)

    if [[ "$HTTP_RES" == *"SUCCEEDED"* ]] || [[ "$HTTP_RES" == *"Resetting"* ]] || [[ "$HTTP_RES" == *"shutdown"* ]]; then
        echo "   ✅ HTTP stop 命令已接收，Arthas 开始重置增强并退出。"
        STOP_SENT=1
    fi
fi

# --- 第 2 级: 尝试通过官方客户端 (arthas-client.jar) 向 Telnet 发送 stop ---
if [[ "$STOP_SENT" -ne 1 ]] && command -v java &>/dev/null; then
    CLIENT_JAR="${PACKAGE_ROOT}/arthas-bin/arthas-client.jar"
    if [[ -f "${CLIENT_JAR}" ]]; then
        echo ">> [2/3] 尝试通过 arthas-client.jar 连接 Telnet 端口 (${TELNET_PORT}) 发送 stop..."
        CLIENT_OUT=$(java -jar "${CLIENT_JAR}" 127.0.0.1 "${TELNET_PORT}" -c "stop" 2>&1 || true)
        if [[ "$CLIENT_OUT" == *"Resetting"* ]] || [[ "$CLIENT_OUT" == *"shutdown"* ]] || [[ "$CLIENT_OUT" == *"server is going to shutdown"* ]]; then
            echo "   ✅ 已通过 Client 发送退出指令。"
            STOP_SENT=1
        fi
    fi
fi

# --- 第 3 级: 尝试通过 as.sh 附加退出 ---
if [[ "$STOP_SENT" -ne 1 && -n "$TARGET_PID" && -f "${PACKAGE_ROOT}/arthas-bin/as.sh" ]]; then
    echo ">> [3/3] 尝试通过 as.sh 向进程 ${TARGET_PID} 发送 stop..."
    "${PACKAGE_ROOT}/arthas-bin/as.sh" "${TARGET_PID}" -c "stop" 2>&1 >/dev/null || true
fi

# --- 第 4 级: 验证端口释放状态 ---
echo ""
echo ">> 正在验证端口 ${HTTP_PORT} 释放状态..."
sleep 2

PORT_RELEASED=0
for i in {1..4}; do
    # 使用 curl 检测是否仍可连通
    if ! curl --noproxy "*" --connect-timeout 1 -m 2 -s "http://127.0.0.1:${HTTP_PORT}/" &>/dev/null; then
        # 进一步通过端口侦听检测
        REMAINING_PID=$(find_pid_by_port "${HTTP_PORT}")
        if [[ -z "$REMAINING_PID" ]]; then
            PORT_RELEASED=1
            break
        fi
    fi
    sleep 1
done

if [[ "$PORT_RELEASED" -eq 1 ]]; then
    echo "================================================================="
    echo "✅ Arthas MCP 会话已彻底停止！"
    echo "   - 端口 ${HTTP_PORT} 已安全释放"
    echo "   - 目标 JVM 字节码增强已完全清除（无残留性能损耗）"
    echo "================================================================="
    exit 0
else
    echo "================================================================="
    echo "⚠️ 提示: 端口 ${HTTP_PORT} 仍未释放。"
    if [[ -n "$TARGET_PID" ]]; then
        echo "   目标进程 PID: ${TARGET_PID}"
        echo "   建议操作:"
        echo "   1. 检查目标 JVM 是否仍在运行；"
        echo "   2. 在连接的 Agent/客户端中发送一次 'stop' 工具调用；"
        echo "   3. 或手动执行: java -jar ${PACKAGE_ROOT}/arthas-bin/arthas-client.jar 127.0.0.1 ${TELNET_PORT} -c 'stop'"
    fi
    echo "================================================================="
    exit 1
fi