#!/bin/bash

# --- 配置区 ---
REMOTE_USER="root"
REMOTE_SERVER="172.24.9.97"
PYTHON_SCRIPT_FILE="./sso_script.py"
# --- 配置区结束 ---

# 检查本地脚本文件是否存在
if [ ! -f "$PYTHON_SCRIPT_FILE" ]; then
    echo "错误: Python脚本文件 $PYTHON_SCRIPT_FILE 不存在!"
    exit 1
fi

echo "第一步: 正在将脚本上传到远程服务器..."

# 在远程创建一个唯一的临时文件名
# 使用 mktemp 在本地生成一个有意义的名字，方便后续引用
# 格式: /tmp/interactive_script_随机数.py
TEMP_SCRIPT_NAME_ON_REMOTE="/srv/sso/cgi-bin/interactive_script_$$.py"

# 使用第一个SSH连接，将本地脚本内容写入远程的临时文件
# 这个连接不需要 -t，因为它只负责传输数据
# 注意：我们使用 `printf` 来构建远程命令，以避免复杂的引号问题
ssh "$REMOTE_USER@$REMOTE_SERVER" "
    CT_ID=\$(pct list | awk '/celery/ {print \$1; exit}');
    if [ -z \"\$CT_ID\" ]; then
        echo '错误: 未找到名称包含 celery 的容器' >&2;
        exit 1;
    fi;
    echo \"找到容器 ID: \$CT_ID\" >&2;
    pct exec \"\$CT_ID\" -- bash -c \"cat > $TEMP_SCRIPT_NAME_ON_REMOTE\"
" < "$PYTHON_SCRIPT_FILE"

# 检查上传是否成功
if [ $? -ne 0 ]; then
    echo "错误: 脚本上传失败。"
    exit 1
fi

echo "上传成功: $TEMP_SCRIPT_NAME_ON_REMOTE"

# 设置本地清理机制，确保即使脚本被中断，也能清理远程临时文件
cleanup_remote() {
    echo "正在清理远程临时文件..." >&2
    ssh "$REMOTE_USER@$REMOTE_SERVER" "
        CT_ID=\$(pct list | awk '/celery/ {print \$1; exit}');
        if [ -n \"\$CT_ID\" ]; then
            pct exec \"\$CT_ID\" -- bash -c \"rm -f $TEMP_SCRIPT_NAME_ON_REMOTE\" 2>/dev/null || true;
        fi
    " 2>/dev/null || true
}
trap cleanup_remote EXIT INT TERM

echo "----------------------------------------------------"
echo "第二步: 正在连接到远程服务器以交互式执行脚本..."
echo "(执行完毕后，临时文件将被自动删除)"
echo "----------------------------------------------------"

# 使用第二个SSH连接，这个连接是纯粹的交互式连接 (-t)
# 它负责执行远程的脚本，并连接你的终端
# 注意：我们使用 `pct exec` 而不是 `pct enter`，这样更脚本化
ssh -t "$REMOTE_USER@$REMOTE_SERVER" "
    CT_ID=\$(pct list | awk '/celery/ {print \$1; exit}');
    if [ -z \"\$CT_ID\" ]; then
        echo '错误: 未找到名称包含 celery 的容器' >&2;
        exit 1;
    fi;
    # 定义一个退出清理函数
    cleanup() {
        echo \"清理远程临时文件...\";
        pct exec \"\$CT_ID\" -- bash -c \"rm -f $TEMP_SCRIPT_NAME_ON_REMOTE\";
        echo \"清理完成。\";
    }
    # 设置脚本结束时（无论是正常退出还是被中断）都执行清理
    trap cleanup EXIT

    # 执行脚本（使用 set -e 确保任何命令失败都会退出）
    pct exec \"\$CT_ID\" -- bash -c \"
        set -e;
        cd /srv/sso/cgi-bin/;
        source /srv/venv/bin/activate;
        python $TEMP_SCRIPT_NAME_ON_REMOTE
    \" || { echo '错误: Python 脚本执行失败' >&2; exit 1; };
"

echo "----------------------------------------------------"
echo "远程交互式脚本已结束。"