#!/bin/sh

set -e

check_command() {
    if [ $? -ne 0 ]; then
        echo "Error: $1 failed"
        exit 1
    fi
}

cd /workspace || {
    echo "Error: Cannot change to /workspace directory"
    exit 1
}

echo "Cloning unsloth-train-scripts repository..."
git clone https://git.hye.su/mira/unsloth-train-scripts.git
check_command "Git clone"

echo "Downloading and running unsloth auto-install script..."
INSTALL_CMD=$(wget -qO- https://raw.githubusercontent.com/unslothai/unsloth/main/unsloth/_auto_install.py | python -)
check_command "Auto-install script"

echo "Executing installation command..."
eval "$INSTALL_CMD"
check_command "Installation command"

echo "Installing additional dependencies..."
pip install gdown wandb
check_command "pip install"

echo "Downloading dataset"
gdown --fuzzy 'https://drive.google.com/file/d/1mqhq69dsSOK7ep7trTjRd3FMagTFTrzF/view?usp=sharing'
check_command "gdown"

echo "Enter your Weights & Biases API key when prompted"
wandb login
check_command "wandb login"

echo "Setup completed successfully."