#!/bin/sh

set -eu

# constants
WORKSPACE_DIR="/workspace"
REPO_URL="https://git.hye.su/mira/unsloth-train-scripts.git"
INSTALL_SCRIPT_URL="https://raw.githubusercontent.com/unslothai/unsloth/main/unsloth/_auto_install.py"
DATASET_URL="https://drive.google.com/file/d/1mqhq69dsSOK7ep7trTjRd3FMagTFTrzF/view?usp=sharing"

error_exit() {
    echo "Error: $1" >&2
    exit 1
}

run_command() {
    if ! "$@"; then
        error_exit "Command failed: $*"
    fi
}

check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        error_exit "This script must be run as root"
    fi
}

setup_workspace() {
    cd "$WORKSPACE_DIR" || error_exit "Cannot change to $WORKSPACE_DIR directory"
}

install_system_packages() {
    echo "Installing system packages..."
    run_command apt-get update
    run_command apt-get install -y cmake tmux vim
}

clone_repository() {
    echo "Cloning unsloth-train-scripts repository..."
    run_command git clone "$REPO_URL"
}

install_unsloth() {
    echo "Downloading and running unsloth auto-install script..."
    INSTALL_CMD=$(wget -qO- "$INSTALL_SCRIPT_URL" | python -) || error_exit "Failed to get install command"
    [ -n "$INSTALL_CMD" ] || error_exit "Empty install command received"
    run_command eval "$INSTALL_CMD"
}

install_python_deps() {
    echo "Installing additional dependencies..."
    run_command pip install gdown wandb
}

download_dataset() {
    echo "Downloading dataset..."
    run_command gdown --fuzzy "$DATASET_URL"
}

setup_wandb() {
    echo "Enter your Weights & Biases API key when prompted"
    run_command wandb login
}

# Main execution
main() {
    check_root
    setup_workspace
    install_system_packages
    clone_repository
    install_unsloth
    install_python_deps
    download_dataset
    setup_wandb
    echo "Setup completed successfully."
}

main