#!/bin/bash
set -e

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

print_header "Starting DreamOmni2 Deployment on Remote Server"

# 1. Check for GPU
print_header "Checking GPU Status"
if ! command -v nvidia-smi &> /dev/null; then
    echo "Error: nvidia-smi not found. Please ensure NVIDIA drivers are installed."
    exit 1
fi
nvidia-smi

# 2. Install Miniconda
print_header "Checking/Installing Miniconda"
if ! command -v conda &> /dev/null; then
    if [ -d "$HOME/miniconda3" ]; then
        echo "Miniconda directory found, adding to path..."
        source "$HOME/miniconda3/bin/activate"
    else
        echo "Miniconda not found. Installing..."
        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
else
    echo "Conda is already installed."
fi

# Ensure conda is available in this shell
eval "$(conda shell.bash hook)"

# 3. Setup Workspace
print_header "Setting up Workspace"
mkdir -p ~/ai-projects
cd ~/ai-projects

# 4. Clone Repository
if [ ! -d "DreamOmni2" ]; then
    echo "Cloning DreamOmni2 repository..."
    git clone https://github.com/dvlab-research/DreamOmni2.git
else
    echo "Repository already exists. Updating..."
    cd DreamOmni2
    git pull
    cd ..
fi

cd DreamOmni2

# 5. Create/Update Conda Environment
print_header "Setting up Conda Environment"
if conda info --envs | grep -q "dreamomni2"; then
    echo "Environment 'dreamomni2' already exists."
else
    echo "Creating environment 'dreamomni2'..."
    conda create -n dreamomni2 python=3.10 -y
fi

echo "Activating environment..."
conda activate dreamomni2

# 6. Install Dependencies
print_header "Installing Python Dependencies"
echo "Installing PyTorch (CUDA 12.1)..."
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

echo "Installing requirements..."
if [ -f "requirements.txt" ]; then
    pip install -r requirements.txt
else
    echo "Warning: requirements.txt not found. Installing basic dependencies..."
    pip install transformers diffusers accelerate sentencepiece protobuf
fi

# Install additional likely required packages
pip install huggingface_hub modelscope

# 7. Download Model Weights
print_header "Downloading Model Weights"
echo "Downloading DreamOmni2 weights from Hugging Face..."
huggingface-cli download xiabs/DreamOmni2 --local-dir checkpoints/DreamOmni2 --local-dir-use-symlinks False

print_header "Deployment Complete!"
echo "Project location: ~/ai-projects/DreamOmni2"
echo "To use the model:"
echo "  1. ssh into the server"
echo "  2. cd ~/ai-projects/DreamOmni2"
echo "  3. conda activate dreamomni2"
echo "  4. Run your inference scripts"