#!/bin/bash
set -e

PROJECT_ROOT="$HOME/projects/web-dev-swarm"
VENV_PATH="$PROJECT_ROOT/venv"

echo "Starting environment setup..."

# 1. Create venv
if [ ! -d "$VENV_PATH" ]; then
    echo "Creating virtual environment in $VENV_PATH..."
    python3 -m venv "$VENV_PATH"
else
    echo "Virtual environment already exists."
fi

# 2. Upgrade pip
echo "Upgrading pip..."
"$VENV_PATH/bin/pip" install --upgrade pip

# 3. Install dependencies
# We'll install them one by one to track progress and avoid massive single-command failures
echo "Installing torch, torchvision, torchaudio..."
"$VENV_PATH/bin/pip" install torch torchvision torchaudio

echo "Installing unsloth and related tools..."
# Using the recommended unsloth install pattern for local
"$VENV_PATH/bin/pip" install unsloth

echo "Installing additional ML tools (bitsandbytes, trl, peft, accelerate, xformers)..."
"$VENV_PATH/bin/pip" install bitsandbytes trl peft accelerate xformers

echo "Verifying installation..."
"$VENV_PATH/bin/python3" -c "import torch; import unsloth; print(f'Torch version: {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}'); print('Unsloth imported successfully!')"

echo "Setup complete!"