setup

#!/bin/sh

# Check if we're using apt
APT="$(which apt)"

install_package() {
    # $@ = packages

    if [ "$APT" = "" ]; then
        echo "This is not an apt system, so cannot install packages. Exiting..."
        exit 1
    fi

    if [ "$(id -u)" -ne 0 ]; then
        NEEDRESTART_MODE=a DEBIAN_FRONTEND=noninteractive sudo apt update
        NEEDRESTART_MODE=a DEBIAN_FRONTEND=noninteractive sudo apt install -y $@
    else
        NEEDRESTART_MODE=a DEBIAN_FRONTEND=noninteractive apt update
        NEEDRESTART_MODE=a DEBIAN_FRONTEND=noninteractive apt install -y $@
    fi
}

ensure_installed() {
    # $1 = command to check
    # $2 = package name

    if ! command -v $1 &> /dev/null
    then
        echo "$1 could not be found, installing..."
        install_package $2
    fi
}

# Let's do some sanity checks. Make sure $HOME is set
if [ -z "$HOME" ]; then
    echo "HOME is not set. Exiting..."
    exit 1
fi

# Go there
GIT_PARENT_DIR="$HOME"
cd "$GIT_PARENT_DIR"

# Set git creds
git config --global credential.helper store
git config --global user.name apaz-cli
git config --global user.email aarpazdera@gmail.com

# Clone the repos in the background while packages are installing
git clone https://github.com/apaz-cli/Dotfiles --depth=1 &
pid1=$!
git clone https://github.com/apaz-cli/Scripts --depth=1 &
pid2=$!

#if [ "$(id -u)" -ne 0 ]; then
#    sudo sed -i "/#\$nrconf{restart} = 'i';/s/.*/\$nrconf{restart} = 'a';/" /etc/needrestart/needrestart.conf
#else
#    sed -i "/#\$nrconf{restart} = 'i';/s/.*/\$nrconf{restart} = 'a';/" /etc/needrestart/needrestart.conf
#fi

# Install the important packages
install_package \
  python3-pip python3-venv \
  nano tree less file locate jq ranger \
  kitty-terminfo \
  git wget curl gcc g++ clang \
  htop nvtop btop \
  numactl libnuma-dev \
  linux-tools-$(uname -r) linux-tools-generic # for perf

# Make sure python aliases python3, and create it if it doesn't
if ! command -v python &> /dev/null
then
    echo "Python alias not found. Creating..."
    PY3PATH="$(which python3)"
    PY_ALIAS_PATH="$(echo $PY3PATH | sed 's/.$//' )" # Remove the last character
    ln -s $PY3PATH $PY_ALIAS_PATH 2> /dev/null
fi

# Move dotfiles into place.
echo "Waiting for repos to clone..."
wait $pid1
wait $pid2
#wait $pid3

# Install secrets repo.
git clone https://github.com/apaz-cli/Secrets --depth=1

echo "Repos are cloned. Moving dotfiles into place..."
cp Dotfiles/.bashrc "$HOME/.bashrc"
cp -r Dotfiles/.nano* "$HOME/"

# Set up prime intellect server
if [ ! "$PRIME" = "" ]; then

    echo "Cloning related repos."
    git clone https://github.com/primeintellect-ai/CPUOptimizer
    git clone https://github.com/primeintellect-ai/pccl

    echo "Setting up Prime..."
    curl -sSL https://raw.githubusercontent.com/PrimeIntellect-ai/prime/main/scripts/install/install.sh > /tmp/prime_install.sh
    chmod +x /tmp/prime_install.sh
    NEEDRESTART_MODE=a DEBIAN_FRONTEND=noninteractive /tmp/prime_install.sh
    rm /tmp/prime_install.sh

    cd prime/
    . .venv/bin/activate
    export LD_LIBRARY_PATH=$(python -c "import site; print(site.getsitepackages()[0] + '/nvidia/nvjitlink/lib')"):$LD_LIBRARY_PATH
fi

echo "\n\e[35mSetup complete.\e[0m\n"