#!/usr/bin/env bash
# Setup script for LMFlow + VILA environment
# This script creates a conda environment with all dependencies needed for training and inference

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Configuration
ENV_FILE="${SCRIPT_DIR}/.dvlm_env_name"
ENV_NAME="${1:-}"

if [ -z "$ENV_NAME" ] && [ -f "$ENV_FILE" ]; then
    ENV_NAME="$(cat "$ENV_FILE")"
fi

if [ -z "$ENV_NAME" ]; then
    read -p "Enter conda env name [dvlm-shiyi]: " ENV_NAME
    ENV_NAME=${ENV_NAME:-"dvlm-shiyi"}
    echo "$ENV_NAME" > "$ENV_FILE"
    echo "[INFO] Saved env name to $ENV_FILE"
else
    echo "$ENV_NAME" > "$ENV_FILE"
    echo "[INFO] Saved env name to $ENV_FILE"
fi

echo "============================================"
echo "LMFlow + VILA Environment Setup"
echo "============================================"
echo ""
echo "Environment name: $ENV_NAME"
echo "Script directory: $SCRIPT_DIR"
echo ""

# Check if conda is available
if ! command -v conda &> /dev/null; then
    echo "[ERROR] conda not found. Please install Anaconda or Miniconda first."
    exit 1
fi

# Initialize conda for bash
eval "$(conda shell.bash hook)"

# Check if environment already exists
if conda env list | grep -q "^${ENV_NAME} "; then
    echo "[WARNING] Environment '$ENV_NAME' already exists."
    read -p "Do you want to remove it and create a new one? (y/N): " confirm
    if [[ $confirm == [yY] || $confirm == [yY][eE][sS] ]]; then
        echo "[INFO] Removing existing environment..."
        conda env remove -n $ENV_NAME -y
    else
        echo "[INFO] Skipping environment creation. Activating existing environment..."
        conda activate $ENV_NAME
    fi
fi

# Create conda environment if it doesn't exist
if ! conda env list | grep -q "^${ENV_NAME} "; then
    echo "[INFO] Creating conda environment from environment.yml..."
    
    # Check if environment.yml exists
    if [ ! -f "$SCRIPT_DIR/environment.yml" ]; then
        echo "[ERROR] environment.yml not found in $SCRIPT_DIR"
        exit 1
    fi
    
    # Create environment
    conda env create -f "$SCRIPT_DIR/environment.yml" -n $ENV_NAME
fi

# Activate environment
echo "[INFO] Activating environment..."
conda activate $ENV_NAME

# Verify Python version
echo "[INFO] Python version: $(python --version)"
echo "[INFO] Python path: $(which python)"

# Install Python dependencies
echo ""
echo "[INFO] Installing Python dependencies from requirements.txt..."
pip install -r "$SCRIPT_DIR/requirements.txt"

# Install Flash Attention 2 (requires special handling)
echo ""
echo "[INFO] Installing Flash Attention 2..."
echo "Note: This may take several minutes to compile."

if python -c "import flash_attn" 2>/dev/null; then
    echo "[INFO] Flash Attention already installed. Skipping."
else
    # Try to install flash-attn from wheel first, fall back to source
    FLASH_ATTN_VERSION="2.5.8"
    CUDA_VERSION=$(python -c "import torch; print(torch.version.cuda.replace('.', '')[:3])" 2>/dev/null || echo "124")
    TORCH_VERSION=$(python -c "import torch; print('.'.join(torch.__version__.split('.')[:2]))" 2>/dev/null || echo "2.6")
    PYTHON_VERSION=$(python -c "import sys; print(f'{sys.version_info.major}{sys.version_info.minor}')" 2>/dev/null || echo "310")

    echo "[INFO] Detected CUDA version: $CUDA_VERSION, PyTorch version: $TORCH_VERSION, Python version: $PYTHON_VERSION"

    # Try wheel installation first
    WHEEL_URL="https://github.com/Dao-AILab/flash-attention/releases/download/v${FLASH_ATTN_VERSION}/flash_attn-${FLASH_ATTN_VERSION}+cu${CUDA_VERSION}torch${TORCH_VERSION}cxx11abiFALSE-cp${PYTHON_VERSION}-cp${PYTHON_VERSION}-linux_x86_64.whl"

    if pip install "$WHEEL_URL" 2>/dev/null; then
        echo "[INFO] Flash Attention installed from pre-built wheel."
    else
        echo "[INFO] Pre-built wheel not available, installing from source..."
        echo "[INFO] This may take 10-30 minutes depending on your system."
        pip install flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation
    fi
fi

# Verify flash-attn installation
if python -c "import flash_attn; print(f'Flash Attention version: {flash_attn.__version__}')" 2>/dev/null; then
    echo "[INFO] Flash Attention installed successfully."
else
    echo "[WARNING] Flash Attention installation may have issues. Training might still work without it."
fi

# Install VILA in editable mode
echo ""
echo "[INFO] Installing VILA..."
if [ -d "$SCRIPT_DIR/VILA" ]; then
    pip install -e "$SCRIPT_DIR/VILA[train,eval]"
    echo "[INFO] VILA installed successfully."
    
    # Copy deepspeed replacement files if they exist
    if [ -d "$SCRIPT_DIR/VILA/llava/train/deepspeed_replace" ]; then
        echo "[INFO] Copying DeepSpeed replacement files..."
        SITE_PACKAGES=$(python -c 'import site; print(site.getsitepackages()[0])')
        cp -rv "$SCRIPT_DIR/VILA/llava/train/deepspeed_replace/"* "$SITE_PACKAGES/deepspeed/" 2>/dev/null || true
    fi
else
    echo "[WARNING] VILA directory not found at $SCRIPT_DIR/VILA"
    echo "[INFO] Skipping VILA installation. You can install it later with:"
    echo "       git clone https://github.com/NVlabs/VILA.git"
    echo "       pip install -e VILA[train,eval]"
fi

# Install LMFlow in editable mode
echo ""
echo "[INFO] Installing LMFlow..."
pip install -e "$SCRIPT_DIR"
echo "[INFO] LMFlow installed successfully."

# Downgrade protobuf for compatibility
echo ""
echo "[INFO] Ensuring protobuf compatibility..."
pip install protobuf==3.20.3

# Final verification
echo ""
echo "============================================"
echo "Installation Complete!"
echo "============================================"
echo ""
echo "Installed packages:"
echo "  - torch: $(python -c 'import torch; print(torch.__version__)')"
echo "  - transformers: $(python -c 'import transformers; print(transformers.__version__)')"
echo "  - deepspeed: $(python -c 'import deepspeed; print(deepspeed.__version__)')"
echo "  - accelerate: $(python -c 'import accelerate; print(accelerate.__version__)')"
python -c "import flash_attn; print(f'  - flash-attn: {flash_attn.__version__}')" 2>/dev/null || echo "  - flash-attn: not installed"
python -c "import vila; print(f'  - vila: {vila.__version__}')" 2>/dev/null || echo "  - vila: not installed"
echo "  - lmflow: $(python -c 'import lmflow; print(lmflow.__version__)')"
echo ""
echo "To activate the environment:"
echo "  conda activate $ENV_NAME"
echo ""
echo "To verify GPU availability:"
echo "  python -c \"import torch; print(f'CUDA available: {torch.cuda.is_available()}, Devices: {torch.cuda.device_count()}')\""
echo ""

