#!/bin/bash
# 唐卡修复AI模型上传脚本
# 上传模型到Hugging Face
# Developed by Wangchuk Mind

set -e  # 遇到错误立即退出

# ===== 配置变量 =====
REPO_ID="Wangchuk1376/ThangkaModels"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOCAL_DIR="$SCRIPT_DIR"

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

# ===== 打印横幅 =====
print_banner() {
    echo -e "${BLUE}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║                                                              ║"
    echo "║          🎨 唐卡修复AI模型 - Hugging Face上传工具 🎨         ║"
    echo "║                                                              ║"
    echo "║          上传模型到: ${REPO_ID}                               ║"
    echo "║                                                              ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
}

# ===== 检查依赖 =====
check_dependencies() {
    echo -e "${YELLOW}🔍 检查依赖...${NC}"
    
    # 检查Hugging Face CLI
    if ! command -v hf &> /dev/null && ! command -v huggingface-cli &> /dev/null; then
        echo -e "${RED}❌ Hugging Face CLI未安装${NC}"
        echo ""
        echo "请安装Hugging Face CLI:"
        echo ""
        echo "  macOS:"
        echo "    brew install huggingface-cli"
        echo ""
        echo "  Linux/Windows:"
        echo "    pip install huggingface_hub[cli]"
        echo ""
        exit 1
    fi
    
    # 统一使用hf命令
    HF_CMD="hf"
    if ! command -v hf &> /dev/null; then
        HF_CMD="huggingface-cli"
    fi
    
    echo -e "${GREEN}✅ Hugging Face CLI已安装${NC}"
    
    # 检查是否已登录
    if ! $HF_CMD whoami &> /dev/null; then
        echo -e "${RED}❌ 未登录Hugging Face${NC}"
        echo ""
        echo "请先登录:"
        echo "  $HF_CMD login"
        echo ""
        exit 1
    fi
    
    CURRENT_USER=$($HF_CMD whoami 2>&1 | head -n 1)
    echo -e "${GREEN}✅ 已登录为: ${CURRENT_USER}${NC}"
    echo ""
}

# ===== 确认上传 =====
confirm_upload() {
    echo -e "${YELLOW}📋 上传信息确认:${NC}"
    echo "  仓库: ${REPO_ID}"
    echo "  本地目录: ${LOCAL_DIR}"
    echo ""
    
    # 统计文件
    echo "  正在统计文件..."
    TOTAL_FILES=$(find "$LOCAL_DIR" -type f ! -path "*/\.*" ! -name "upload_to_huggingface.sh" | wc -l | tr -d ' ')
    TOTAL_SIZE=$(du -sh "$LOCAL_DIR" | cut -f1)
    
    echo "  文件数量: ${TOTAL_FILES}"
    echo "  总大小: ${TOTAL_SIZE}"
    echo ""
    
    # 列出主要文件
    echo "  主要文件:"
    if [ -f "$LOCAL_DIR/README.md" ]; then
        echo "    ✅ README.md"
    else
        echo "    ⚠️  README.md (未找到)"
    fi
    
    if [ -d "$LOCAL_DIR/models" ]; then
        echo "    ✅ models/ 目录"
        echo "       - sd2.1_base_paddle/"
        echo "       - control_v11p_sd21_canny_paddle/"
        echo "       - finetuned/ (LoRA模型)"
    else
        echo "    ⚠️  models/ 目录 (未找到)"
    fi
    echo ""
    
    # 询问确认
    read -p "确认上传? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        echo -e "${YELLOW}❌ 上传已取消${NC}"
        exit 0
    fi
    echo ""
}

# ===== 创建.gitattributes =====
create_gitattributes() {
    echo -e "${YELLOW}📝 创建.gitattributes...${NC}"
    
    cat > "$LOCAL_DIR/.gitattributes" << 'EOF'
# 使用Git LFS跟踪大文件
*.safetensors filter=lfs diff=lfs merge=lfs -text
*.pdparams filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.tar.gz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
EOF
    
    echo -e "${GREEN}✅ .gitattributes创建完成${NC}"
    echo ""
}

# ===== 上传文件 =====
upload_files() {
    echo -e "${BLUE}📤 开始上传文件...${NC}"
    echo ""
    
    cd "$LOCAL_DIR"
    
    # 步骤1: 上传README和配置文件
    echo -e "${YELLOW}步骤 1/3: 上传文档和配置文件...${NC}"
    if [ -f "README.md" ]; then
        $HF_CMD upload "$REPO_ID" README.md --repo-type model
        echo -e "${GREEN}✅ README.md 上传完成${NC}"
    fi
    
    if [ -f ".gitattributes" ]; then
        $HF_CMD upload "$REPO_ID" .gitattributes --repo-type model
        echo -e "${GREEN}✅ .gitattributes 上传完成${NC}"
    fi
    echo ""
    
    # 步骤2: 上传models目录
    if [ -d "models" ]; then
        echo -e "${YELLOW}步骤 2/3: 上传models目录...${NC}"
        echo "  ⏳ 这可能需要较长时间，请耐心等待..."
        echo ""
        
        $HF_CMD upload "$REPO_ID" models/ \
            --repo-type model \
            --num-workers 4 \
            --verbose
        
        echo ""
        echo -e "${GREEN}✅ models目录上传完成${NC}"
    fi
    echo ""
    
    # 步骤3: 上传其他文件
    echo -e "${YELLOW}步骤 3/3: 上传其他文件...${NC}"
    
    # 查找其他需要上传的文件
    OTHER_FILES=$(find . -maxdepth 1 -type f \
        ! -name "README.md" \
        ! -name ".gitattributes" \
        ! -name ".DS_Store" \
        ! -name "upload_to_huggingface.sh" \
        ! -path "*/.*")
    
    if [ -n "$OTHER_FILES" ]; then
        for file in $OTHER_FILES; do
            filename=$(basename "$file")
            $HF_CMD upload "$REPO_ID" "$file" --repo-type model
            echo -e "${GREEN}✅ ${filename} 上传完成${NC}"
        done
    else
        echo "  没有其他文件需要上传"
    fi
    echo ""
}

# ===== 验证上传 =====
verify_upload() {
    echo -e "${BLUE}🔍 验证上传...${NC}"
    echo ""
    
    REPO_URL="https://huggingface.co/${REPO_ID}"
    
    echo "  模型仓库地址: ${REPO_URL}"
    echo ""
    
    # 尝试打开浏览器
    if command -v open &> /dev/null; then
        read -p "是否在浏览器中打开? (Y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Nn]$ ]]; then
            open "$REPO_URL"
        fi
    fi
    echo ""
}

# ===== 显示完成信息 =====
show_completion() {
    echo -e "${GREEN}"
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║                                                              ║"
    echo "║                    🎉 上传完成！ 🎉                           ║"
    echo "║                                                              ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo -e "${NC}"
    echo ""
    echo "📦 模型仓库: https://huggingface.co/${REPO_ID}"
    echo ""
    echo "📚 使用方法:"
    echo ""
    echo "  # 使用CLI下载"
    echo "  huggingface-cli download ${REPO_ID} --local-dir ./models"
    echo ""
    echo "  # 使用Python下载"
    echo "  from huggingface_hub import snapshot_download"
    echo "  snapshot_download(repo_id=\"${REPO_ID}\", local_dir=\"./models\")"
    echo ""
    echo "  # 在代码中使用"
    echo "  from diffusion_paddle import load_model"
    echo "  pipe = load_model(\"${REPO_ID}/sd2.1_base_paddle\")"
    echo ""
    echo "🌟 别忘了给项目点星！"
    echo "   GitHub: https://github.com/WangchukMind/thangka-restoration-ai"
    echo ""
}

# ===== 主函数 =====
main() {
    print_banner
    
    # 检查依赖
    check_dependencies
    
    # 确认上传
    confirm_upload
    
    # 创建.gitattributes
    create_gitattributes
    
    # 上传文件
    upload_files
    
    # 验证上传
    verify_upload
    
    # 显示完成信息
    show_completion
}

# ===== 运行主函数 =====
main "$@"