#!/usr/bin/bash

script_version=0.3.1
echo "Script Version: $script_version"
#default stuff
username=""
renderkey=""
autosettings=""
gpu_cores=""
client_name=""
run_cpu_client="n"
bypass_confirm="n"

while [[ $# -gt 0 ]]; do
    key="$1"
    case $key in
        -username)
            username="$2"
        shift 2;;
        -renderkey)
            renderkey="$2"
        shift 2;;
        -autosettings)
            autosettings="$2"
        shift 2;;
        -gpucores)
            gpu_cores="$2"
        shift 2;;
        -clientname)
            client_name="$2"
        shift 2;;
        -cores)
            cpu_cores="$2"
        shift 2;;
        -cpu)
            if [ "$2" == "y" ]; then
                run_cpu_client="y"
            else
                run_cpu_client="n"
            fi
        shift 2;;
        -y)
            bypass_confirm="y"
        shift 1;;
        *)
            echo "Unknown option: $1"
            exit 1
        ;;
    esac
done
line_thing="------------------------------------"
amount_of_gpus=$(nvidia-smi -L | wc -l)
amount_of_ram=$(free -h | awk '/^Mem:/ {print $7}' | sed 's/[^0-9]*//g')
cpu_cores_fixed=$(nproc)

echo "$line_thing"
echo "Welcome to the Vast.ai Sheepit setup script! Please follow the on-screen instructions to continue."
echo "$line_thing"

if [ -z "$username" ]; then
    read -r -p "Enter your username: " username < /dev/tty
fi
if [ -z "$renderkey" ]; then
    read -r -p "Enter your render key: " renderkey < /dev/tty
fi
if [ -z "$autosettings" ]; then
    read -r -p "Do you want to use auto settings? (y/n): " autosettings < /dev/tty
fi
if [ -z "$client_name" ]; then
    read -r -p "Enter client name (no spaces): " client_name < /dev/tty
fi


if [[ "$autosettings" =~ ^[Yy]$ ]]; then
    echo "$line_thing"
    echo "Auto settings enabled."
    num_gpus=$amount_of_gpus
    gpu_rec_amount_ram=$(( amount_of_ram / num_gpus ))
    RAM=$gpu_rec_amount_ram
    if [ -z "$gpu_cores" ]; then
        gpu_cores=2
    fi
else
    echo "$line_thing"
    echo "There are $amount_of_gpus NVIDIA GPUs in this machine."
    while true; do
        read -r -p "Enter the number of GPUs you would like to use for sheepit (1-$amount_of_gpus): " num_gpus < /dev/tty
        if [[ "$num_gpus" =~ ^[1-9][0-9]*$ && "$num_gpus" -le "$amount_of_gpus" ]]; then
            break
        else
            echo "Please enter a valid positive integer for the number of GPUs (1-$amount_of_gpus)."
        fi
    done
    gpu_rec_amount_ram=$(( amount_of_ram / num_gpus ))
    amount_of_gpus1=$(( num_gpus + 1 ))
    cpu_rec_amount_ram=$(( amount_of_ram / amount_of_gpus1 ))
    echo "$line_thing"
    echo "This machine has ${amount_of_ram}GB RAM available."
    echo "It is recommended to use ${gpu_rec_amount_ram}GB of RAM per GPU client, or ${cpu_rec_amount_ram}GB if also running a CPU client."
    read -r -p "Enter amount of RAM for each client in GB: " RAM < /dev/tty
    RAM="${RAM//[^0-9]/}"
    while true; do
        read -r -p "Enter number of cores for each GPU client: " gpu_cores < /dev/tty
        if [[ "$gpu_cores" =~ ^[1-9][0-9]*$ ]]; then
            break
        else
            echo "Please enter a valid positive integer."
        fi
    done
fi

if [[ "$run_cpu_client" == "y" ]]; then
    WANT_CPU="y"
else
    if [[ "$autosettings" =~ ^[Yy]$ ]]; then
        WANT_CPU="n"
    else
        read -r -p "Do you want to run a CPU client as well? (y/n): " WANT_CPU < /dev/tty
    fi
fi

if [[ "$WANT_CPU" =~ ^[Yy]$ ]]; then
    cores_used_for_gpu=$(( num_gpus * gpu_cores ))
    cores_left=$(( cpu_cores_fixed - cores_used_for_gpu ))
    if [[ "$autosettings" =~ ^[Yy]$ ]]; then
        cpu_cores=$cores_left
    else
        
        echo "This machine has $cpu_cores_fixed CPU cores."
        echo "With GPU clients using $cores_used_for_gpu cores, there are $cores_left cores available for the CPU client."
        while true; do
            read -r -p "Enter number of cores for the CPU client: " cpu_cores < /dev/tty
            if [[ "$cpu_cores" =~ ^[1-9][0-9]*$ ]]; then
                break
            else
                echo "Please enter a valid positive integer."
            fi
        done
    fi
fi

echo "$line_thing"
echo "These are the settings selected, is this ok? (Y/n)"
echo "Client name: $client_name"
echo "GPU Clients: $num_gpus"
echo "RAM per client: ${RAM}GB"
echo "Cores per GPU: ${gpu_cores}"
echo "Sheepit user: ${username}"
echo "Renderkey: ${renderkey}"
if [[ "$WANT_CPU" =~ ^[Yy]$ ]]; then
    echo "Run a CPU client: Yes"
    echo "CPU cores: $cpu_cores"
else
    echo "Run a CPU client: No"
fi

if [[ "$bypass_confirm" != "y" ]]; then
    read -r -p "(Y/n): " confirm < /dev/tty
    if [[ "$confirm" =~ ^[Nn]$ ]]; then
        echo "Aborted"
        exit 0
    fi
fi

echo "Starting sheepit..."

apt update
apt install -y sudo
sudo apt update && sudo apt upgrade -y
sudo apt install -y tmux default-jre blender wget xserver-xorg-video-dummy libboost-all-dev libgl1-mesa-dev libglu1-mesa libsm-dev libxkbcommon0 btop
sudo mkdir -p /sheepit/sheepitcache
sudo chown -R "$(whoami):$(id -gn)" /sheepit
wget https://www.sheepit-renderfarm.com/media/applet/client-latest.php -O /sheepit/client.jar
tmux new-session -d -s sheepit

for (( i = 0; i < num_gpus; i++ )); do
    GPU_ID="OPTIX_$i"
    CACHE_DIR="/sheepit/sheepitcache/$i"
    HOSTNAME="${client_name}-$(printf "%03d" $(( i + 1 )))"
    
    CMD="java -jar /sheepit/client.jar \
    -login $username -password $renderkey \
    -ui text --verbose --headless -cache-dir $CACHE_DIR \
    -gpu $GPU_ID -cores $gpu_cores -memory ${RAM}G -hostname $HOSTNAME -shared-zip /sheepit/sheepitcache/shared"
    
    if [ "$i" -eq 0 ]; then
        tmux rename-window -t sheepit "gpu-$i"
        tmux send-keys -t sheepit "$CMD" C-m
    else
        tmux new-window -t sheepit -n "gpu-$i"
        tmux send-keys -t sheepit:"gpu-$i" "$CMD" C-m
    fi
done

if [[ "$WANT_CPU" =~ ^[Yy]$ ]]; then
    CPU_CACHE_DIR="/sheepit/sheepitcache/cpu"
    CPU_HOSTNAME="${client_name}-cpu"
    
    CPU_CMD="java -jar /sheepit/client.jar \
    -login $username -password $renderkey \
    -ui text --verbose --headless -cache-dir $CPU_CACHE_DIR \
    -cores $cpu_cores -memory ${RAM}G -hostname $CPU_HOSTNAME -shared-zip /sheepit/sheepitcache/shared"
    
    tmux new-window -t sheepit -n "cpu"
    tmux send-keys -t sheepit:"cpu" "$CPU_CMD" C-m
fi

# Attach to the tmux session
tmux switch-client -t sheepit
tmux attach-session -t sheepit
echo "If it doesn't switch automatically, you can switch with 'tmux attach-session -t sheepit' or 'tmux switch-client -t sheepit'"