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

# 防止 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 credential.helper store
echo "https://Shuyun18:${HF_TOKEN}@huggingface.co" > /root/.git-credentials

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

elif [ "$1" == "save" ]; then
    echo "Saving data to dataset..."
    if [ ! -d "$LOCAL_REPO/.git" ]; then
        mkdir -p "$LOCAL_REPO"
        cd "$LOCAL_REPO"
        git init
        git remote add origin "$REPO_URL"
    else
        cd "$LOCAL_REPO"
    fi
    
    # 确保本地和目标子目录存在
    mkdir -p "$LOCAL_REPO/hermes" "$LOCAL_REPO/open-webui"
    mkdir -p "$HERMES_DIR" "$OPENWEBUI_DIR"
    
    rsync -av --delete --exclude='.git' "$HERMES_DIR/" "$LOCAL_REPO/hermes/"
    rsync -av --delete --exclude='.git' "$OPENWEBUI_DIR/" "$LOCAL_REPO/open-webui/"
    
    git add .
    git commit -m "Auto-sync data at $(date)" || echo "Nothing to commit"
    git push origin main || echo "Push failed"
    echo "Data saved."
fi