#!/bin/bash

SOURCE_ROOT="/root/comfyui/ComfyUI"
DEST_ROOT="/home/waas/ComfyUI"
DATASET_SRC="/datasets/ComfyUI/models"
SPECIAL_PATH="Joy_caption_two/text_model" # 需要强制复制的特殊路径

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

status() {
  case $1 in
    error)   echo -e "${RED}[ERROR]${RESET} $2" ;;
    success) echo -e "${GREEN}[SUCCESS]${RESET} $2" ;;
    warn)    echo -e "${YELLOW}[WARNING]${RESET} $2" ;;
    info)    echo -e "${BLUE}[INFO]${RESET} $2" ;;
  esac
}

initialize_comfyui() {

    
  # 处理主目录结构
  if [ ! -d "$DEST_ROOT" ]; then
    status info "初始化ComfyUI目录结构..."
    
    # 创建目标目录
    mkdir -pv "$DEST_ROOT" || {
      status error "创建目标目录失败"
      return 1
    }

   

    # 移动并创建软链接
    for dir in models output; do
      if [ -d "${SOURCE_ROOT}/${dir}" ]; then
        # 移动目录
        mv -v "${SOURCE_ROOT}/${dir}" "$DEST_ROOT/" || {
          status error "移动 ${dir} 目录失败"
          return 1
        }
        
        # 创建软链接
        ln -svf "${DEST_ROOT}/${dir}" "${SOURCE_ROOT}/${dir}" || {
          status error "创建 ${dir} 软链接失败"
          return 1
        }
      else
        status warn "源目录 ${SOURCE_ROOT}/${dir} 不存在，跳过移动"
        # 仍然创建目标
        mkdir -p "${DEST_ROOT}/${dir}"
      fi
    done
    
  else
    status info "检测到已存在的ComfyUI目录，更新软链接..."
    
    # 不论是否存在目标目录都创建
    for dir in output models ; do
        mkdir -p "${DEST_ROOT}/${dir}" || {
            status info "目标文件夹${DEST_ROOT}/${dir}已存在，创建失败"
        }
    done
   
    # 删除源目录并创建软链接
    for dir in output models ; do
      if [ -e "${SOURCE_ROOT}/${dir}" ]; then

        if [ -L "${SOURCE_ROOT}/${dir}" ]; then
            
          status info "软链接 ${SOURCE_ROOT}/${dir} 已存在，跳过"
          continue
        fi

        
        # 删除实体目录
        rm -rfv "${SOURCE_ROOT}/${dir}" || {
          status error "删除 ${dir} 目录失败"
          return 1
        }
      fi
      
      # 创建新软链接
      ln -svf "${DEST_ROOT}/${dir}" "${SOURCE_ROOT}/${dir}" || {
        status error "创建 ${dir} 软链接失败"
        return 1
      }
    done
  fi
}

process_special_directory() {
  local src_fullpath="${DATASET_SRC}/${SPECIAL_PATH}"
  local dest_fullpath="${DEST_ROOT}/models/${SPECIAL_PATH}"

  mkdir -pv "${dest_fullpath%/*}" || {
    status error "无法创建父目录 ${dest_fullpath%/*}"
    return 1
  }

  if [ -d "$src_fullpath" ]; then
    status info "强制复制特殊目录: ${SPECIAL_PATH}"
    cp -nrv "$src_fullpath" "${dest_fullpath%/*}/" 2>&1 | while read -r line; do
      status info "COPY: $line"
    done
  else
    status warn "源特殊目录不存在: $src_fullpath"
  fi
}

process_datasets() {
  status info "开始处理数据集"
  
  process_special_directory

  find "${DATASET_SRC}" -mindepth 1 -maxdepth 2 -type d -print0 | while IFS= read -r -d $'\0' dir; do
    relative_path="${dir#$DATASET_SRC/}"
    [[ "$relative_path" == "$SPECIAL_PATH" ]] && continue
    
    dest_dir="${DEST_ROOT}/models/${relative_path}"
    
    mkdir -pv "$dest_dir" || continue
    
    find "$dir" -mindepth 1 -maxdepth 1 -print0 | while IFS= read -r -d $'\0' item; do
      item_name=$(basename "$item")
      [ ! -e "$dest_dir/$item_name" ] && ln -svf "$item" "$dest_dir/"
    done
  done
}

# 主执行流程
if initialize_comfyui; then
  process_datasets
  status success "操作完成"
else
  status error "初始化失败，中止执行"
  exit 1
fi