#!/bin/bash

# GitHub Copilot 远程开发连接设置脚本
# 作者: zhch158
# 版本: 2.2

# ==================== 配置参数 ====================
LOCAL_PROXY_PORT=7890
REMOTE_FORWARD_PORT=7281
REMOTE_HOST="10.192.72.11"
REMOTE_USER="ubuntu"
VLLM_PORT=8101
LOCAL_VLLM_PORT=8111

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

# ==================== 工具函数 ====================
log_info() {
    echo -e "${BLUE}ℹ${NC} $1"
}

log_success() {
    echo -e "${GREEN}✅${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}⚠${NC} $1"
}

log_error() {
    echo -e "${RED}❌${NC} $1"
}

log_step() {
    echo -e "${PURPLE}🔧${NC} $1"
}

# ==================== 检查函数 ====================
check_dependencies() {
    log_step "检查依赖工具..."
    
    local missing_tools=()
    
    # 检查必需工具
    for tool in ssh curl netstat; do
        if ! command -v $tool >/dev/null 2>&1; then
            missing_tools+=($tool)
        fi
    done
    
    if [ ${#missing_tools[@]} -ne 0 ]; then
        log_error "缺少必需工具: ${missing_tools[*]}"
        log_info "请安装缺少的工具后重试"
        exit 1
    fi
    
    log_success "所有依赖工具已安装"
}

check_local_proxy() {
    log_step "检查本地代理服务..."
    
    # 检查端口是否监听
    if netstat -tln 2>/dev/null | grep -q ":$LOCAL_PROXY_PORT "; then
        log_info "发现本地端口 $LOCAL_PROXY_PORT 正在监听"
    else
        log_warning "本地端口 $LOCAL_PROXY_PORT 未在监听"
    fi
    
    # 测试代理连接
    if curl -s --proxy http://localhost:$LOCAL_PROXY_PORT --connect-timeout 5 http://www.google.com > /dev/null 2>&1; then
        log_success "本地代理 (端口 $LOCAL_PROXY_PORT) 工作正常"
        return 0
    else
        log_warning "本地代理不可用，将使用直接连接"
        return 1
    fi
}

check_remote_connectivity() {
    log_step "检查远程主机连接..."
    
    if ssh -o ConnectTimeout=10 -o BatchMode=yes $REMOTE_USER@$REMOTE_HOST "echo 'SSH连接成功'" 2>/dev/null; then
        log_success "远程主机 $REMOTE_HOST 连接正常"
        return 0
    else
        log_error "无法连接到远程主机 $REMOTE_HOST"
        log_info "请检查网络连接和SSH配置"
        return 1
    fi
}

# ==================== 主要功能 ====================
setup_ssh_tunnel() {
    local proxy_target
    
    if check_local_proxy; then
        proxy_target="localhost:$LOCAL_PROXY_PORT"
    else
        # 如果没有本地代理，使用一个无效的目标让连接失败
        log_warning "没有可用的本地代理，GitHub Copilot可能无法正常工作"
        proxy_target="127.0.0.1:1"  # 故意使用一个无效端口
    fi
    
    log_step "建立SSH隧道连接..."
    echo -e "${CYAN}📡 端口转发配置:${NC}"
    echo "   本地 :$LOCAL_VLLM_PORT -> 远程 :$VLLM_PORT (vLLM 服务)"
    echo "   远程 :$REMOTE_FORWARD_PORT -> 本地 :$LOCAL_PROXY_PORT (代理服务)"
    echo ""
    
    # SSH 连接参数优化，减少警告
    ssh \
        -o "ExitOnForwardFailure=yes" \
        -o "ServerAliveInterval=60" \
        -o "ServerAliveCountMax=3" \
        -o "TCPKeepAlive=yes" \
        -o "Compression=yes" \
        -o "StrictHostKeyChecking=no" \
        -o "UserKnownHostsFile=/dev/null" \
        -o "LogLevel=ERROR" \
        -L $LOCAL_VLLM_PORT:localhost:$VLLM_PORT \
        -R $REMOTE_FORWARD_PORT:$proxy_target \
        $REMOTE_USER@$REMOTE_HOST \
        "$(cat << 'REMOTE_SCRIPT'
#!/bin/bash

# 远程执行脚本
REMOTE_FORWARD_PORT=7281
SHELL_DIR=$HOME/zhch/shell
mkdir -p $SHELL_DIR

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
NC='\033[0m'

log_info() { echo -e "${BLUE}ℹ${NC} $1"; }
log_success() { echo -e "${GREEN}✅${NC} $1"; }
log_warning() { echo -e "${YELLOW}⚠${NC} $1"; }
log_error() { echo -e "${RED}❌${NC} $1"; }
log_step() { echo -e "${PURPLE}🔧${NC} $1"; }

echo -e "${CYAN}🚀 在远程主机上配置 GitHub Copilot 环境${NC}"
echo "=================================================="

# 备份原有代理设置
log_step "备份原有代理设置..."
env | grep -i proxy > $SHELL_DIR/copilot_proxy_backup 2>/dev/null || true

# 显示当前环境
log_info "当前代理环境:"
env | grep -i proxy | sed 's/^/   /' || echo "   (无代理设置)"

# 清除现有代理设置
log_step "清除现有代理设置..."
unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
unset ftp_proxy FTP_PROXY no_proxy NO_PROXY

# 设置新的代理环境变量
log_step "设置隧道代理环境变量..."
export http_proxy=http://localhost:$REMOTE_FORWARD_PORT
export https_proxy=http://localhost:$REMOTE_FORWARD_PORT
export HTTP_PROXY=http://localhost:$REMOTE_FORWARD_PORT
export HTTPS_PROXY=http://localhost:$REMOTE_FORWARD_PORT
export no_proxy=localhost,127.0.0.1,10.192.72.11,172.16.0.0/12
export NO_PROXY=localhost,127.0.0.1,10.192.72.11,172.16.0.0/12

log_success "代理环境变量已设置"
echo "   http_proxy=http://localhost:$REMOTE_FORWARD_PORT"
echo "   https_proxy=http://localhost:$REMOTE_FORWARD_PORT"

# 检查端口监听状态
log_step "检查端口转发状态..."
sleep 2  # 等待端口转发建立

if netstat -tln 2>/dev/null | grep -q "127.0.0.1:$REMOTE_FORWARD_PORT " || ss -tln 2>/dev/null | grep -q ":$REMOTE_FORWARD_PORT "; then
    log_success "端口 $REMOTE_FORWARD_PORT 正在监听"
else
    log_warning "端口 $REMOTE_FORWARD_PORT 未检测到监听状态"
fi

# 测试网络连接
log_step "测试网络连接..."

echo "🧪 连接测试结果:"

# 测试隧道代理
if curl -s --proxy http://localhost:$REMOTE_FORWARD_PORT --connect-timeout 10 http://httpbin.org/ip > /dev/null 2>&1; then
    log_success "隧道代理连接正常"
else
    log_warning "隧道代理连接失败"
fi

# 测试 GitHub API
if curl -s --connect-timeout 10 https://api.github.com/user > /dev/null 2>&1; then
    log_success "GitHub API 可访问"
else
    log_error "GitHub API 不可访问"
fi

# 测试 GitHub Copilot 服务
if curl -s --connect-timeout 10 https://copilot-proxy.githubusercontent.com/ > /dev/null 2>&1; then
    log_success "GitHub Copilot 服务可访问"
else
    log_warning "GitHub Copilot 服务访问可能有问题"
fi

# 显示环境信息
echo ""
echo -e "${CYAN}📊 环境信息摘要${NC}"
echo "=================================================="
echo "🔗 SSH 隧道: 已建立"
echo "🌐 代理端口: localhost:$REMOTE_FORWARD_PORT"
echo "📝 备份文件: $SHELL_DIR/copilot_proxy_backup"
echo "💻 主机名: $(hostname)"
echo "👤 用户: $(whoami)"
echo "📅 时间: $(date)"

# 创建快速测试脚本
cat > $SHELL_DIR/test_copilot_connection.sh << 'TEST_EOF'
#!/bin/bash
echo "🧪 GitHub Copilot 连接测试"
echo "=========================="

# 测试代理连接
if curl -s --connect-timeout 5 https://api.github.com > /dev/null; then
    echo "✅ GitHub API: 可访问"
else
    echo "❌ GitHub API: 不可访问"
fi

if curl -s --connect-timeout 5 https://copilot-proxy.githubusercontent.com/ > /dev/null; then
    echo "✅ Copilot Service: 可访问"
else
    echo "❌ Copilot Service: 不可访问"
fi

echo ""
echo "📊 当前代理设置:"
env | grep -i proxy | sort
TEST_EOF

chmod +x $SHELL_DIR/test_copilot_connection.sh

# 创建代理恢复脚本
cat > $SHELL_DIR/restore_proxy.sh << 'RESTORE_EOF'
#!/bin/bash
echo "🔄 恢复原有代理设置..."

if [ -f $SHELL_DIR/copilot_proxy_backup ]; then
    source $SHELL_DIR/copilot_proxy_backup
    echo "✅ 代理设置已恢复"
    echo "📊 当前代理设置:"
    env | grep -i proxy | sort
else
    echo "❌ 未找到备份文件"
fi
RESTORE_EOF

chmod +x $SHELL_DIR/restore_proxy.sh

echo ""
echo -e "${GREEN}🎉 GitHub Copilot 环境配置完成!${NC}"
echo "=================================================="
echo "📝 使用说明:"
echo "   • 保持此 SSH 会话活跃以维持 Copilot 功能"
echo "   • 在 VS Code 中正常使用 GitHub Copilot"
echo "   • 运行 $SHELL_DIR/test_copilot_connection.sh 测试连接"
echo "   • 运行 $SHELL_DIR/restore_proxy.sh 恢复原有代理设置"
echo ""
echo "🔧 调试命令:"
echo "   netstat -tln | grep 127.0.0.1:$REMOTE_FORWARD_PORT  # 检查端口"
echo "   curl -v https://api.github.com             # 测试 GitHub API"
echo "   env | grep -i proxy                        # 查看代理设置"
echo ""

# 设置退出时恢复代理的陷阱
trap 'echo "🔄 会话结束，恢复代理设置..."; source $SHELL_DIR/copilot_proxy_backup 2>/dev/null || true; echo "✅ 代理设置已恢复"' EXIT

log_info "保持 SSH 会话活跃中... (按 Ctrl+C 退出)"

# 保持会话活跃并定期显示状态
while true; do
    sleep 60
    echo "📡 $(date '+%H:%M:%S'): SSH 隧道活跃，Copilot 代理运行中"
    
    # 检查代理连接状态
    if netstat -an 2>/dev/null | grep -q "127.0.0.1:$REMOTE_FORWARD_PORT.*ESTABLISHED"; then
        echo "🔥 检测到活跃的代理连接"
    fi
done

'REMOTE_SCRIPT'
)"
}

# ==================== 清理函数 ====================
cleanup() {
    log_info "正在清理..."
    # 这里可以添加清理逻辑
}

# ==================== 主程序 ====================
main() {
    echo -e "${CYAN}🚀 GitHub Copilot 远程开发连接设置${NC}"
    echo "=================================================="
    echo "📅 $(date)"
    echo "🖥️  本机: $(hostname)"
    echo "🌐 远程: $REMOTE_HOST"
    echo ""
    
    # 设置退出陷阱
    trap cleanup EXIT
    
    # 执行检查
    check_dependencies
    check_remote_connectivity || exit 1
    
    # 设置隧道
    setup_ssh_tunnel
}

# ==================== 脚本入口 ====================
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi