#!/bin/bash

echo "=== 配置 Hugging Face SSH 连接 ==="
echo ""

# 1. 检查并创建 SSH 目录
echo "1. 检查 SSH 目录..."
mkdir -p ~/.ssh
chmod 700 ~/.ssh
echo "✓ SSH 目录已准备"
echo ""

# 2. 检查 SSH key
echo "2. 检查 SSH key..."
if [ -f ~/.ssh/id_ed25519.pub ]; then
    echo "✓ 找到 id_ed25519.pub"
    SSH_KEY_FILE="~/.ssh/id_ed25519"
    SSH_PUB_KEY="~/.ssh/id_ed25519.pub"
elif [ -f ~/.ssh/id_rsa.pub ]; then
    echo "✓ 找到 id_rsa.pub"
    SSH_KEY_FILE="~/.ssh/id_rsa"
    SSH_PUB_KEY="~/.ssh/id_rsa.pub"
else
    echo "✗ 未找到 SSH key，正在生成..."
    ssh-keygen -t ed25519 -f ~/.ssh/id_ed25519 -N "" -C "huggingface"
    SSH_KEY_FILE="~/.ssh/id_ed25519"
    SSH_PUB_KEY="~/.ssh/id_ed25519.pub"
    echo "✓ 已生成新的 SSH key"
fi
echo ""

# 3. 显示公钥
echo "3. 你的 SSH 公钥（请复制并添加到 Hugging Face）："
echo "----------------------------------------"
if [ -f ~/.ssh/id_ed25519.pub ]; then
    cat ~/.ssh/id_ed25519.pub
elif [ -f ~/.ssh/id_rsa.pub ]; then
    cat ~/.ssh/id_rsa.pub
fi
echo "----------------------------------------"
echo ""
echo "⚠️  请访问 https://huggingface.co/settings/keys 添加上述公钥"
echo "   按 Enter 继续（添加完公钥后）..."
read
echo ""

# 4. 配置 SSH config
echo "4. 配置 SSH config..."
SSH_CONFIG="$HOME/.ssh/config"

# 备份现有配置
if [ -f "$SSH_CONFIG" ]; then
    cp "$SSH_CONFIG" "$SSH_CONFIG.backup.$(date +%Y%m%d_%H%M%S)"
    echo "✓ 已备份现有配置"
fi

# 确定使用的 key 文件
if [ -f ~/.ssh/id_ed25519 ]; then
    KEY_PATH="$HOME/.ssh/id_ed25519"
elif [ -f ~/.ssh/id_rsa ]; then
    KEY_PATH="$HOME/.ssh/id_rsa"
fi

# 添加或更新 hf.co 配置
if grep -q "Host hf.co" "$SSH_CONFIG" 2>/dev/null; then
    echo "✓ 更新现有 hf.co 配置"
    # 删除旧的 hf.co 配置块
    sed -i '/^Host hf.co$/,/^$/d' "$SSH_CONFIG"
fi

# 添加新配置
cat >> "$SSH_CONFIG" << EOF

# Hugging Face SSH 配置
Host hf.co
    HostName hf.co
    User git
    Port 22
    IdentityFile $KEY_PATH
    IdentitiesOnly yes
    StrictHostKeyChecking accept-new
EOF

chmod 600 "$SSH_CONFIG"
echo "✓ SSH config 已配置（端口 22）"
echo ""

# 5. 设置 key 文件权限
echo "5. 设置文件权限..."
chmod 600 "$KEY_PATH"
chmod 644 "$KEY_PATH.pub"
echo "✓ 权限已设置"
echo ""

# 6. 测试 SSH 连接（端口 22）
echo "6. 测试 SSH 连接（端口 22）..."
if ssh -T git@hf.co 2>&1 | grep -q "successfully authenticated"; then
    echo "✓ SSH 连接成功！"
    echo ""
    echo "=== 配置完成 ==="
    echo "现在可以正常使用 git push 了"
    exit 0
else
    echo "✗ 端口 22 连接失败，尝试配置端口 443..."
    echo ""
    
    # 更新配置为端口 443
    sed -i 's/Port 22/Port 443/' "$SSH_CONFIG"
    echo "✓ 已更新为端口 443"
    echo ""
    
    echo "7. 测试 SSH 连接（端口 443）..."
    if ssh -T git@hf.co 2>&1 | grep -q "successfully authenticated"; then
        echo "✓ SSH 连接成功（使用端口 443）！"
        echo ""
        echo "=== 配置完成 ==="
        echo "现在可以正常使用 git push 了"
        exit 0
    else
        echo "✗ 端口 443 也连接失败"
        echo ""
        echo "请检查："
        echo "1. SSH 公钥是否已添加到 https://huggingface.co/settings/keys"
        echo "2. 网络连接是否正常"
        echo "3. 防火墙是否阻止了连接"
        echo ""
        echo "可以手动测试："
        echo "  ssh -T git@hf.co"
        exit 1
    fi
fi