#!/bin/bash

#SIA_INSTALL_NO_WEB=1
SIA_INSTALL_NO_NOTEBOOK=1
#SIA_INSTALL_NO_CORE=1
#SIA_INSTALL_NO_ITB=1
SIA_INSTALL_NO_MISTRAL_INFER=1
#SIA_INSTALL_NO_GEMMA_INFER=1
#SIA_INSTALL_NO_GEMMA_TRAIN=1

trap 'return 1' ERR

# Check if required environment variables are set
if [ -z "$SIA_REPO_URL" ]; then
    echo "Error: SIA_REPO_URL environment variable is not set"
    return 1
fi

if [ -z "$SIA_REPO_USER" ]; then
    echo "Error: SIA_REPO_USER environment variable is not set"
    return 1
fi

if [ -z "$SIA_REPO_PAT" ]; then
    echo "Error: SIA_REPO_PAT environment variable is not set"
    return 1
fi

# Install required packages
apt-get update
apt-get install -y \
    build-essential \
    cmake \
    cuda-toolkit \
    curl \
    git \
    gnupg \
    jq \
    libcurl4-nss-dev \
    python3-dev \
    python3-venv \
    tmux \
    vim \
    wget \
    ;

# Install Chrome
if [ -z "${SIA_INSTALL_NO_ITB}" ]; then
    wget -q -O - https://dl-ssl.google.com/linux/linux_signing_key.pub | gpg --dearmor -o /etc/apt/trusted.gpg.d/google.gpg
    echo "deb http://dl.google.com/linux/chrome/deb/ stable main" >> /etc/apt/sources.list.d/google-chrome.list
    apt-get update && \
        apt-get install -y \
        google-chrome-stable
fi

# Install Node.js
if [ -z "${SIA_INSTALL_NO_CORE}" ]; then
    echo "Installing Node.js..."
    curl -fsSL https://deb.nodesource.com/setup_20.x | bash -
    apt-get install -y nodejs
fi

# Install Rust
curl https://sh.rustup.rs -sSf | bash -s -- -y
export PATH="/root/.cargo/bin:${PATH}"

# Install llama.cpp
cd /
curl -O https://git.nielsgeens.be/api/packages/llm/generic/llama.cpp/b5269/llama.cpp.tar
tar -xf /llama.cpp.tar -C /usr/local/bin --wildcards --no-anchored "llama-*"
tar -xf /llama.cpp.tar -C /usr/local/lib --wildcards --no-anchored "*.so"
rm llama.cpp.tar

# Create directory structure
echo "Creating directory structure..."
mkdir -p "/root/data/iterations"
mkdir -p "/root/desktop"
mkdir -p "/root/venvs"

# Clone SIA repository
if [ ! -d "/root/sia" ]; then
    echo "Cloning SIA repository..."
    git clone --recursive "https://$SIA_REPO_USER:$SIA_REPO_PAT@${SIA_REPO_URL#https://}" "/root/sia"
    
    echo "Configuring git user and email..."
    cd "/root/sia"
    git config --global user.name "Niels Geens"
    git config --global user.email "niels.geens@gmail.com"
    git config --global core.editor vim
fi

# Install SIA dependencies
cd /root/desktop
source "/root/sia/scripts/install.sh"

# Finetune model
if [ -z "${SIA_INSTALL_NO_TRAIN}" ]; then
    echo "Finetuning model..."
    /root/sia/scripts/train.sh /root/models/bootstrap
    ln -s /root/models/bootstrap /root/models/current
fi

# Start SIA using restart script
if [ -z "${SIA_INSTALL_NO_CORE}" ]; then
    echo "Run restart script..."
    "/root/sia/scripts/restart.sh"
fi