#!/usr/bin/env bash
# 断点续传大文件到云服务器 /home/ubuntu/data
# 用法: bash scripts/sync_data_to_server.sh [all|weights|sd|hairstyles]
set -euo pipefail

REMOTE="ubuntu@117.50.213.111"
REMOTE_BASE="/home/ubuntu/data"
LOCAL_BASE="/home/xsl/change_hair"
LOG_DIR="${LOCAL_BASE}/project/logs"
mkdir -p "$LOG_DIR"

RSYNC_OPTS=(-avP --partial --info=progress2)

sync_weights() {
  echo ">>> [1/3] 同步 hair_service_sd/weights (~10G) ..."
  rsync "${RSYNC_OPTS[@]}" \
    "${LOCAL_BASE}/project/hair_service_sd/weights/" \
    "${REMOTE}:${REMOTE_BASE}/hair_service_sd/weights/"
}

sync_sd_models() {
  echo ">>> [2/3] 同步 SD 底模 + ESRGAN + GFPGAN (~4.6G) ..."
  rsync "${RSYNC_OPTS[@]}" \
    "${LOCAL_BASE}/project/onediff/stable-diffusion-webui/models/Stable-diffusion/v1-5-pruned-emaonly.safetensors" \
    "${REMOTE}:${REMOTE_BASE}/sdwebui/models/Stable-diffusion/"
  rsync "${RSYNC_OPTS[@]}" \
    "${LOCAL_BASE}/project/onediff/stable-diffusion-webui/models/ESRGAN/" \
    "${REMOTE}:${REMOTE_BASE}/sdwebui/models/ESRGAN/"
  rsync "${RSYNC_OPTS[@]}" \
    "${LOCAL_BASE}/project/onediff/stable-diffusion-webui/models/GFPGAN/" \
    "${REMOTE}:${REMOTE_BASE}/sdwebui/models/GFPGAN/"
}

sync_hairstyles() {
  echo ">>> [3/3] 同步 5 个 chang_* 发型资源 (~0.8G) ..."
  for id in chang_tuoyuan chang_bolang chang_zhixian chang_huaban chang_xinxing; do
    echo "  - $id"
    rsync "${RSYNC_OPTS[@]}" \
      "${LOCAL_BASE}/project/data/ref_hairstyle/${id}/" \
      "${REMOTE}:${REMOTE_BASE}/project/data/ref_hairstyle/${id}/"
    rsync "${RSYNC_OPTS[@]}" \
      "${LOCAL_BASE}/project/data/hair_template_material/${id}/" \
      "${REMOTE}:${REMOTE_BASE}/project/data/hair_template_material/${id}/"
    rsync "${RSYNC_OPTS[@]}" \
      "${LOCAL_BASE}/project/data/upload_train_imgs/${id}/" \
      "${REMOTE}:${REMOTE_BASE}/project/data/upload_train_imgs/${id}/"
    rsync "${RSYNC_OPTS[@]}" \
      "${LOCAL_BASE}/data/train_material/${id}/" \
      "${REMOTE}:${REMOTE_BASE}/train_material/${id}/"
  done
}

TARGET="${1:-all}"
case "$TARGET" in
  weights)   sync_weights ;;
  sd)        sync_sd_models ;;
  hairstyles) sync_hairstyles ;;
  all)
    sync_weights
    sync_sd_models
    sync_hairstyles
    ;;
  *)
    echo "用法: $0 [all|weights|sd|hairstyles]"
    exit 1
    ;;
esac

echo ">>> 同步完成: ${TARGET}"