#!/bin/bash
set -e

# Function to print section headers
print_header() {
    echo "=================================================="
    echo "$1"
    echo "=================================================="
}

print_header "Configuring DreamOmni2 on Remote Server"

# 1. Check for GPU
if ! command -v nvidia-smi &> /dev/null; then
    echo "Error: nvidia-smi not found."
    exit 1
fi

# 2. Install Miniconda if needed
if ! command -v conda &> /dev/null; then
    if [ -d "$HOME/miniconda3" ]; then
        source "$HOME/miniconda3/bin/activate"
    else
        echo "Installing Miniconda..."
        mkdir -p ~/miniconda3
        # 使用更稳定的清华源或者官方源下载
        wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
        bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
        rm ~/miniconda3/miniconda.sh
        source ~/miniconda3/bin/activate
        conda init --all
    fi
fi
eval "$(conda shell.bash hook)"

# 修复 Conda Terms of Service 报错
print_header "Checking Conda Configuration"
# 尝试接受 TOS（如果命令存在）
if conda --help | grep -q "tos"; then
    echo "Accepting Conda Terms of Service..."
    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main || true
    conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r || true
fi
# 设置通道优先级，避免一些问题
conda config --set channel_priority flexible || true

# 3. Setup Environment
mkdir -p ~/ai-projects
cd ~/ai-projects

# Check if DreamOmni2 exists (handling potential failed previous uploads)
if [ ! -d "DreamOmni2" ] && [ ! -d "/root/DreamOmni2" ]; then
     echo "Warning: DreamOmni2 directory not found in ~/ai-projects. Checking root..."
     if [ -d "/root/DreamOmni2" ]; then
        mv /root/DreamOmni2 .
     else
        echo "Error: DreamOmni2 code not found."
        exit 1
     fi
fi

cd DreamOmni2

# Check for Model
if [ ! -d "checkpoints/DreamOmni2" ]; then
    echo "Creating checkpoints directory..."
    mkdir -p checkpoints/DreamOmni2
    
    # Check possible upload locations
    MODEL_SOURCE=""
    if [ -d "../../DreamOmni2_Model" ]; then
        MODEL_SOURCE="../../DreamOmni2_Model"
    elif [ -d "/root/DreamOmni2_Model" ]; then
        MODEL_SOURCE="/root/DreamOmni2_Model"
    fi

    if [ -n "$MODEL_SOURCE" ]; then
        echo "Moving uploaded model files from $MODEL_SOURCE..."
        mv "$MODEL_SOURCE"/* checkpoints/DreamOmni2/
        # Remove empty dir
        rmdir "$MODEL_SOURCE" || true
    else
        echo "Warning: Model files not found yet. Please ensure they are uploaded."
    fi
fi

print_header "Setting up Conda Environment"
if conda info --envs | grep -q "dreamomni2"; then
    echo "Environment 'dreamomni2' already exists."
else
    # 使用 conda-forge 可能更稳定，或者尝试默认
    echo "Creating environment..."
    conda create -n dreamomni2 python=3.10 -y
fi
conda activate dreamomni2

print_header "Installing Dependencies"
# 安装核心依赖
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers diffusers accelerate sentencepiece protobuf huggingface_hub modelscope

print_header "Setup Complete!"
echo "To run the model:"
echo "1. ssh -p 5204 root@156.225.27.52"
echo "2. cd ~/ai-projects/DreamOmni2"
echo "3. conda activate dreamomni2"