#!/bin/bash
#export WELCOME_TITLE="Welcome to My SerVice"
#grep "$WELCOME_TITLE" /root/.bashrc &> /dev/null || echo 'echo -e "\033[1;32m'${WELCOME_TITLE}'\033[0m"' >> /root/.bashrc
#!/bin/bash
[ -f /etc/storage/file/base-storage ]  && rm /etc/storage/file/base-storage
[ -f /var/log/storage_error.log ]  && rm /var/log/storage_error.log

Comfy_UI_Migrate3() {
SOURCE_ROOT="/root/comfyui/ComfyUI/"
DEST_ROOT="/home/waas/ComfyUI/"

if [ ! -d ${DEST_ROOT} ];then
        mkdir -pv $DEST_ROOT
        cp -r ${SOURCE_ROOT}/models ${SOURCE_ROOT}/output  $DEST_ROOT
else
    [ ! -d ${DEST_ROOT}/models ] &&   cp -r ${SOURCE_ROOT}/models $DEST_ROOT
    [ ! -d ${DEST_ROOT}/output ] &&   cp -r ${SOURCE_ROOT}/output $DEST_ROOT
    if [ -d ${DEST_ROOT}/models  ];then
	    SOURCE_DIR=$(find "${SOURCE_ROOT}/models" -mindepth 1 -maxdepth 1 -type d -exec basename {} \;)
	    DEST_DIR=$(find "${DEST_ROOT}/models" -mindepth 1 -maxdepth 1 -type d -exec basename {} \;)
	    DIFF_DIR=$(comm -23 <(echo "$SOURCE_DIR" | sort) <(echo "$DEST_DIR" | sort))
	    #echo $DIFF_DIR
	    if [[ ! -z ${DIFF_DIR} ]] ;then
		    cd $SOURCE_ROOT/models && cp -r  $DIFF_DIR $DEST_ROOT/models
	    fi
    fi
fi

#if [ ! -d  /home/waas/ComfyUI/models/LLM/Florence-2-base ];then
#	cp -r /datasets/ComfyUI/models/LLM/Florence-2-base/ /home/waas/ComfyUI/models/LLM/Florence-2-base
#fi
#修改程序的用户存放模型文件路径
sed -i.bak -e 's/^models_dir = .*/models_dir = \"\/home\/waas\/ComfyUI\/models\"/' -e 's/^output_directory = .*/output_directory = \"\/home\/waas\/ComfyUI\/output\"/'   /root/comfyui/ComfyUI/folder_paths.py


cat > /root/comfyui/ComfyUI/extra_model_paths.yaml <<"EOF"
#Rename this to extra_model_paths.yaml and ComfyUI will load it

#config for a1111 ui
#all you have to do is change the base_path to where yours is installed
a111:
    base_path: /datasets/ComfyUI/

    checkpoints: models/checkpoints
    configs: models/configs
    vae: models/vae
    loras: |
         models/loras
         models/LyCORIS
    upscale_models: |
                  models/ESRGAN
                  models/RealESRGAN
                  models/SwinIR
    embeddings: embeddings
    hypernetworks: models/hypernetworks
    controlnet: models/ControlNet
    grounding-dino: models/grounding-dino
    sams: models/sams
    style_models: models/style_models
    clip: models/clip
    clip_vision: models/clip_vision
    unet: models/unet
    Joy_caption_two: models/Joy_caption_two
    LLM: models/LLM
    text_encoders: models/text_encoders
    dz_facedetailer: models/dz_facedetailer
    controlnet: models/controlnet
    upscale_models: models/upscale_models
    Janus-Pro: models/Janus-Pro


#config for comfyui
#your base path should be either an existing comfy install or a central folder where you store all of your models, loras, etc.

#comfyui:
#     base_path: path/to/comfyui/
#     checkpoints: models/checkpoints/
#     clip: models/clip/
#     clip_vision: models/clip_vision/
#     configs: models/configs/
#     controlnet: models/controlnet/
#     embeddings: models/embeddings/
#     loras: models/loras/
#     upscale_models: models/upscale_models/
#     vae: models/vae/

#other_ui:
#    base_path: path/to/ui
#    checkpoints: models/checkpoints
#    gligen: models/gligen
#    custom_nodes: path/custom_nodes
EOF

}
Comfy_UI_Migrate4() {
SOURCE_ROOT="/root/comfyui/ComfyUI/"
DEST_ROOT="/home/waas/ComfyUI/"

if [ ! -d ${DEST_ROOT} ];then
        mkdir -pv $DEST_ROOT
        cp -r ${SOURCE_ROOT}/models ${SOURCE_ROOT}/output  $DEST_ROOT
else
    [ ! -d ${DEST_ROOT}/models ] &&   cp -r ${SOURCE_ROOT}/models $DEST_ROOT
    [ ! -d ${DEST_ROOT}/output ] &&   cp -r ${SOURCE_ROOT}/output $DEST_ROOT
    if [ -d ${DEST_ROOT}/models  ];then
	    SOURCE_DIR=$(find "${SOURCE_ROOT}/models" -mindepth 1 -maxdepth 1 -type d -exec basename {} \;)
	    DEST_DIR=$(find "${DEST_ROOT}/models" -mindepth 1 -maxdepth 1 -type d -exec basename {} \;)
	    DIFF_DIR=$(comm -23 <(echo "$SOURCE_DIR" | sort) <(echo "$DEST_DIR" | sort))
	    #echo $DIFF_DIR
	    if [[ ! -z ${DIFF_DIR} ]] ;then
		    cd $SOURCE_ROOT/models && cp -r  $DIFF_DIR $DEST_ROOT/models
	    fi
    fi
fi

if [ ! -d  /home/waas/ComfyUI/models/LLM/Florence-2-base ];then
	cp -r /datasets/ComfyUI/models/LLM/Florence-2-base/ /home/waas/ComfyUI/models/LLM/Florence-2-base
fi
#修改程序的用户存放模型文件路径
sed -i.bak -e 's/^models_dir = .*/models_dir = \"\/home\/waas\/ComfyUI\/models\"/' -e 's/^output_directory = .*/output_directory = \"\/home\/waas\/ComfyUI\/output\"/'   /root/comfyui/ComfyUI/folder_paths.py


cat > /root/comfyui/ComfyUI/extra_model_paths.yaml <<"EOF"
#Rename this to extra_model_paths.yaml and ComfyUI will load it

#config for a1111 ui
#all you have to do is change the base_path to where yours is installed
a111:
    base_path: /datasets/ComfyUI/

    checkpoints: models/checkpoints
    configs: models/configs
    vae: models/vae
    loras: |
         models/loras
         models/LyCORIS
    upscale_models: |
                  models/ESRGAN
                  models/RealESRGAN
                  models/SwinIR
    embeddings: embeddings
    hypernetworks: models/hypernetworks
    controlnet: models/ControlNet
    grounding-dino: models/grounding-dino
    sams: models/sams
    style_models: models/style_models
    clip: models/clip
    clip_vision: models/clip_vision
    unet: models/unet
    Joy_caption_two: models/Joy_caption_two
    LLM: models/LLM
    text_encoders: models/text_encoders
    dz_facedetailer: models/dz_facedetailer
    controlnet: models/controlnet
    upscale_models: models/upscale_models
    Janus-Pro: models/Janus-Pro
    text_encoders: models/text_encoders
    diffusion_models: models/diffusion_models
    diffusers: models/diffusers
    insightface: models/insightface
    ipadapter-flux: models/ipadapter-flux
    liveportrait: models/liveportrait
    pulid: models/pulid
    segformer_b3_clothes: models/segformer_b3_clothes
    swap_model: models/swap_model
    ipadapter: models/ipadapter
    inisghtface: models/inisghtface


#config for comfyui
#your base path should be either an existing comfy install or a central folder where you store all of your models, loras, etc.

#comfyui:
#     base_path: path/to/comfyui/
#     checkpoints: models/checkpoints/
#     clip: models/clip/
#     clip_vision: models/clip_vision/
#     configs: models/configs/
#     controlnet: models/controlnet/
#     embeddings: models/embeddings/
#     loras: models/loras/
#     upscale_models: models/upscale_models/
#     vae: models/vae/

#other_ui:
#    base_path: path/to/ui
#    checkpoints: models/checkpoints
#    gligen: models/gligen
#    custom_nodes: path/custom_nodes
EOF

}
Echomimic_init() {

# 定义路径变量
DEST_DIR="/home/waas/Echomimic-V2"
MODELS_DIR="$DEST_DIR/models"
OUTPUTS_SRC="/root/echomimic_v2/outputs"
OUTPUTS_DST="$DEST_DIR/outputs"
OUTPUTS_LINK="/root/echomimic_v2/outputs"

# 创建目标目录和models子目录
mkdir -p "$MODELS_DIR"

# 移动outputs文件夹到新的位置

if [ ! -d /home/waas/Echomimic-V2/outputs  ];then  
	mv "$OUTPUTS_SRC" "$DEST_DIR"
fi

# 检查移动操作是否成功
if [ $? -eq 0 ]; then
    echo "outputs 文件夹已成功移动到 $DEST_DIR."

    # 创建软连接回到原来的位置
    ln -sfn "$OUTPUTS_DST" "$OUTPUTS_LINK"
    
    if [ $? -eq 0 ]; then
        echo "创建软连接成功: $OUTPUTS_LINK -> $OUTPUTS_DST."
    else
        echo "创建软连接失败!"
        exit 1
    fi
else
    echo "移动 outputs 文件夹失败!"
    exit 1
fi

echo "脚本执行完毕."
}
WebUI_init() {
# 定义路径变量
DEST_DIR="/home/waas/SDWebui/"
MODELS_SOURCE="/root/stable-diffusion-webui-1.9.4/models"
OUTPUTS_DST="$DEST_DIR/outputs"
#OUTPUTS_LINK="/root/echomimic_v2/outputs"
EXTENSION_SRC="/root/stable-diffusion-webui-1.9.4/extensions"

# 创建目标目录和models子目录
if [ ! -d ${DEST_DIR} ] ;then
        mkdir -p $DEST_DIR
        mkdir -p $OUTPUTS_DST
        mv $EXTENSION_SRC $DEST_DIR
        mv $MODELS_SOURCE $DEST_DIR
else
        if [ ! -d $OUTPUTS_DST ];then
                mkdir -p $OUTPUTS_DST
        fi
        if [ ! -d $DEST_DIR/models ];then
                mv $MODELS_SOURCE $DEST_DIR
        fi
        if [ ! -d $DEST_DIR/extensions ];then
                mv $EXTENSION_SRC $DEST_DIR
        fi

fi
# 移动outputs文件夹到新的位置
}

## 新增函数 加入websocketd监听日志服务
Websocket_log() {
    local port="$1"
    local logPath="$2"
    # 参数完整性检查
    if [[ -z "$port" || -z "$logPath" ]]; then
        echo "ERROR: 参数不能为空" >&2
        return 1
    fi
    echo "监听websocketd日志 端口[$port] 日志路径[$logPath]"
    /usr/bin/websocketd --port "$port" tail -f "$logPath" &

}

websocket_service() {
    local port="$1"
    local cmd="$2"

    if [[ -z "$port" || -z "$cmd" ]]; then
        echo "ERROR: 参数不能为空" >&2
        return 1
    fi
    echo "启动websocketd服务 端口[$port] 命令行[$cmd]"
    /usr/bin/websocketd --port "$port" $cmd &
}

Websocket_ComfyUI_tar() {
    local port="$1"
    local path="/home/waas/ComfyUI/"
    mkdir -p "$path"
    /usr/bin/websocketd --port "$port" bash -c "
        tarName=\"output_\$(date +%Y%m%d-%H%M%S).tar\"          
        tarPath=\"${path}\${tarName}\"                           
        tar -cvf \"\$tarPath\" -C \"${path}\" output &&          
        echo \"文件已打包，请前往文件存储下 ${path} 下载 \${tarName}\"
    " &
}

# WAAS_WEBSOCKETD_PORT_0=6666
# WAAS_WEBSOCKETD_CMD_0="/root/test.sh"
websocket_service() {
    for var in $(compgen -A variable); do
        # 检查变量名是否符合 WAAS_WEBSOCKETD_PORT_X 的规则
        if [[ "$var" =~ ^WAAS_WEBSOCKETD_PORT_([0-9]+)$ ]]; then
            # 提取 X 的值
            index="${BASH_REMATCH[1]}"
            # 获取对应的 PORT 和 CMD 的值
            port="${!var}"
            cmd_var="WAAS_WEBSOCKETD_CMD_${index}"
            cmd="${!cmd_var}"

            # 检查 CMD 是否存在
            if [[ -z "$cmd" ]]; then
                echo "ERROR: 未找到对应的 CMD 变量: $cmd_var"
                continue
            fi

            echo "启动websocketd服务 端口[$port] 命令行[$cmd]"
            /usr/bin/websocketd --port="$port" bash -c "$cmd" &
        fi
    done
}

mkdir /etc/waas-script
mkdir -p /var/log/waas
wget -q  https://file.aigate.cc/waas/east2/v1/waas-welcome.sh -O /etc/waas-script/waas-welcome.sh && echo "下载waas-welcome.sh文件成功" || echo "下载waas-welcome.sh文件失败"
wget -q  https://file.aigate.cc/waas/east2/v1/cron-root-stoage.sh -O /etc/waas-script/cron-root-stoage.sh && echo "下载cron-root-stoage.sh文件成功" || echo "下载cron-root-stoage.sh文件失败"
wget -q  https://file.aigate.cc/waas/east2/v1/proxy.sh -O /etc/waas-script/proxy.sh && echo "下载proxy.sh文件成功" || echo "下载proxy.sh文件失败"
wget -q  https://file.aigate.cc/waas/east2/v1/unset_proxy.sh -O /etc/waas-script/unset_proxy.sh && echo "下载unset_proxy.sh文件成功" || echo "下载unset_proxy.sh文件失败"
wget -q  https://file.aigate.cc/waas/east2/v1/restart_comfyui.sh -O /etc/waas-script/restart_comfyui.sh && echo "下载restart_comfyui.sh文件成功" || echo "下载restart_comfyui.sh文件失败"
wget -q  http://file.aigate.cc/server_script/proccess-compose/base/process-compose.yml -O /etc/waas-script/process-compose.yml  && echo "下载process-compose.yml文件成功" || echo "下载process-compose.yml文件失败"

chmod +x /etc/waas-script/cron-root-stoage.sh
chmod +x /etc/waas-script/restart_comfyui.sh

## 下载websocketd可执行文件
wget -q https://file.aigate.cc/waas/east2/v2/websocketd -O /usr/bin/websocketd && echo "下载websocketd成功" || echo "下载websocketd失败"
chmod +x /usr/bin/websocketd

grep "source /etc/waas-script/waas-welcome.sh" /root/.bashrc &> /dev/null || echo "source /etc/waas-script/waas-welcome.sh " >> /root/.bashrc
grep "alias df='df /home/waas'" /root/.bashrc &> /dev/null  || echo "alias df='df /home/waas'" >> /root/.bashrc
if printenv |grep CUDA_VISIBLE_DEVICES &> /dev/null ;then
       sed -i '2i export CUDA_VISIBLE_DEVICES= ' /etc/waas-script/waas-welcome.sh
       grep "export CUDA_VISIBLE_DEVICES=" /root/.bashrc &> /dev/null ||  echo "export CUDA_VISIBLE_DEVICES=" >> /root/.bashrc
       grep "alias nvidia-smi='echo -n 不支持'" /root/.bashrc &> /dev/null ||  echo "alias nvidia-smi='echo -n "不支持"'" >> /root/.bashrc
       grep "readonly CUDA_VISIBLE_DEVICES" /root/.bashrc &> /dev/null || echo "readonly CUDA_VISIBLE_DEVICES" >>  /root/.bashrc
else
	grep "export CUDA_VISIBLE_DEVICES=" /etc/waas-script/waas-welcome.sh &> /dev/null && sed -i '/export CUDA_VISIBLE_DEVICES=/d' /etc/waas-script/waas-welcome.sh
	grep "export CUDA_VISIBLE_DEVICES=" /root/.bashrc &> /dev/null &&  sed -i '/export CUDA_VISIBLE_DEVICES=/d' /root/.bashrc
	grep "alias nvidia-smi='echo -n 不支持'" /root/.bashrc &> /dev/null &&  sed -i "/alias nvidia-smi='echo -n 不支持'/d" /root/.bashrc
	grep "readonly CUDA_VISIBLE_DEVICES" /root/.bashrc &> /dev/null && sed -i '/readonly CUDA_VISIBLE_DEVICES/d' /root/.bashrc
fi
	

printenv |grep WAAS_COMFYUI_IMAGE_V3   &> /dev/null  && Comfy_UI_Migrate3
printenv |grep WAAS_COMFYUI_IMAGE_V4   &> /dev/null  && Comfy_UI_Migrate4
printenv |grep WAAS_COMFYUI_IMAGE_V5   &> /dev/null  && wget -qO- https://file.aigate.cc/waas/east2/v1/comfyui_v5.sh  | bash &
printenv |grep ECHOMIMIC_IMAGE   &> /dev/null  &&  Echomimic_init
printenv |grep SDWEBUI_WAAS   &> /dev/null  &&  WebUI_init
if printenv |grep SDWEBUI_WAAS   &> /dev/null ;then
	#grep "export SDWEBUI_WAAS=TRUE" /root/.bashrc &> /dev/null || echo  "export SDWEBUI_WAAS=TRUE" >>  /root/.bashrc
	grep "SDWEBUI_WAAS=TRUE" /etc/waas-script/waas-welcome.sh &> /dev/null ||   sed -i '2i export SDWEBUI_WAAS=TRUE' /etc/waas-script/waas-welcome.sh
fi
if printenv |grep ECHO_BASH   &> /dev/null ;then
   grep ${ECHO_BASH} &> /dev/null || echo $ECHO_BASH >> /root/.bashrc
fi


cat << EOF > /root/miniconda3/share/jupyter/kernels/python3/kernel.json
{
 "argv": [
  "python",
  "-m",
  "ipykernel_launcher",
  "-f",
  "{connection_file}"
 ],
 "display_name": "Python 3 (ipykernel)",
 "language": "python",
 "metadata": {
  "debugger": true
 },
 "env": {
        "PATH": "/root/miniconda3/bin:/root/miniconda3/condabin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
        "SHELL": "/bin/bash"
    }
}
EOF


if printenv |grep SSHD_ONLY_IMAGE   &> /dev/null ;then
   ln -s /usr/local/cuda/bin/* /usr/local/bin/
fi
if [ -z "$WAAS_IMAGE_VERSION" ]; then
  while true;do /etc/waas-script/cron-root-stoage.sh ;sleep 30;done &
  
  printenv |grep SSHD_ONLY_IMAGE &> /dev/null && /usr/sbin/sshd -D $SSHD_OPTS  
  
  printenv |grep SSHD_ONLY_IMAGE &> /dev/null || /usr/sbin/sshd -D $SSHD_OPTS &
  
  printenv |grep SSHD_ONLY_IMAGE &> /dev/null || /root/miniconda3/bin/jupyter-lab --allow-root --ServerApp.token="$PASSWORD" --notebook-dir='/'   &
  if printenv |grep OLLAMA_WAAS_IMAGES_DEEPSEEK_32B &> /dev/null ;then
  	/usr/bin/nohup  /usr/bin/ollama serve > /var/log/ollamaserv.log 2>&1 &
  	/usr/bin/ollama run deepseek-r1:32b > /var/log/32b.log 2>&1 &
  	cd /home/webui/nextjs-ollama-llm-ui && /usr/local/bin/npm run dev  > /var/log/ollama-webui.log 2>&1  &
  fi
fi
## 基于环境变量 判断是否启用对应的日志监听服务
if [[ -v WAAS_LOG_COMFYUI_PORT && -v WAAS_LOG_COMFYUI_PATH  ]]; then 
    Websocket_log "$WAAS_LOG_COMFYUI_PORT" "$WAAS_LOG_COMFYUI_PATH"
fi

## 基于环境变量 创建websocketd服务
websocket_service

printenv |grep WAAS_COMFYUI_TAR_PORT   &> /dev/null  && Websocket_ComfyUI_tar "$WAAS_COMFYUI_TAR_PORT"

if [ -z "$WAAS_IMAGE_VERSION" ];then
	printenv |grep SSHD_ONLY_IMAGE &> /dev/null || /usr/bin/code-server --auth password --bind-addr 0.0.0.0:8080  --locale="zh-cn"
fi

if [[ "$IMAGE_TAG" == "T2ITrainer" ]]; then
    wget -q  https://file.aigate.cc/waas/east2/v1/restart_t2itrainer.sh -O /etc/waas-script/restart_t2itrainer.sh && echo "下载restart_t2itrainer.sh文件成功" || echo "下载restart_t2itrainer.sh文件失败"
    chmod +x /etc/waas-script/restart_t2itrainer.sh
    bash -c /etc/waas-script/restart_t2itrainer.sh || \
        echo "restart_t2itrainer.sh 执行失败 - 继续后续操作" >&2
fi

if [[ "$MUSUBI_TUNER" == "TRUE" ]]; then
    echo "musubi-tuner 初始化"
    wget -q  https://file.aigate.cc/waas/east1/v1/musubi-tuner.sh -O /etc/waas-script/musubi-tuner.sh && echo "下载musubi-tuner.sh文件成功" || echo "下载musubi-tuner.sh文件失败"
    chmod +x /etc/waas-script/musubi-tuner.sh
    bash -c /etc/waas-script/musubi-tuner.sh || \
        echo "musubi-tuner.sh 执行失败 - 继续后续操作" >&2
fi

if [[ "$WAAS_IMAGE_VERSION" == "V2" ]]; then
    process-compose --no-server --tui=false up -f /etc/waas-script/process-compose.yml
fi

