#!/bin/bash
REPO_URL="https://huggingface.co/datasets/Shuyun18/hermes"
LOCAL_REPO="/tmp/hermes_data"
HERMES_DIR="/root/.hermes"
HERMES_HUDUI_DIR="/root/.hermes-hudui"

# 防止 Git 因交互式提示挂起
export GIT_TERMINAL_PROMPT=0
export GIT_SSH_COMMAND="ssh -o StrictHostKeyChecking=no"

if [ -z "$HF_TOKEN" ]; then
    echo "Warning: HF_TOKEN not found. Data sync disabled."
    exit 0
fi

# 配置 Git 用户信息
git config --global user.name "Shuyun18"
git config --global user.email "shuyun18@users.noreply.huggingface.co"

# 验证当前认证状态
pip install huggingface_hub -q
echo "=== HF Authentication Status ==="
hf whoami
echo "================================"

# 使用带 token 的 URL 进行 git 操作
REPO_URL_WITH_TOKEN="https://Shuyun18:${HF_TOKEN}@huggingface.co/datasets/Shuyun18/hermes"

if [ "$1" == "load" ]; then
    echo "Loading data from dataset..."
    rm -rf "$LOCAL_REPO"
    git clone "$REPO_URL_WITH_TOKEN" "$LOCAL_REPO" || echo "Dataset repo is empty or not found."
    
    # 确保本地核心目录存在
    mkdir -p "$HERMES_DIR" "$HERMES_HUDUI_DIR"
    
    # 确保从仓库拉取的子目录存在，避免 rsync 报错
    mkdir -p "$LOCAL_REPO/hermes" "$LOCAL_REPO/hermes-hudui"
    
    # 将子目录数据分别同步到对应位置
    rsync -av --exclude='.git' "$LOCAL_REPO/hermes/" "$HERMES_DIR/"
    rsync -av --exclude='.git' "$LOCAL_REPO/hermes-hudui/" "$HERMES_HUDUI_DIR/"
    
    echo "Data loaded."

elif [ "$1" == "save" ]; then
    echo "Saving data to dataset..."
    
    # 先 clone 仓库（确保有正确的 git 历史）
    rm -rf "$LOCAL_REPO"
    git clone "$REPO_URL_WITH_TOKEN" "$LOCAL_REPO" || {
        echo "Failed to clone repo, creating new one..."
        mkdir -p "$LOCAL_REPO"
        cd "$LOCAL_REPO"
        git init
        git checkout -b main
        git remote add origin "$REPO_URL_WITH_TOKEN"
    }
    cd "$LOCAL_REPO"
    
    # 确保本地和目标子目录存在
    mkdir -p "$LOCAL_REPO/hermes" "$LOCAL_REPO/hermes-hudui"
    mkdir -p "$HERMES_DIR" "$HERMES_HUDUI_DIR"
    
    rsync -av --delete --exclude='.git' "$HERMES_DIR/" "$LOCAL_REPO/hermes/"
    rsync -av --delete --exclude='.git' "$HERMES_HUDUI_DIR/" "$LOCAL_REPO/hermes-hudui/"
    
    git add .
    git commit -m "Auto-sync data at $(date)" || echo "Nothing to commit"
    git push -u origin main || echo "Push failed"
    echo "Data saved."
fi