#!/bin/bash

# Set the directory for WEBUI
WEBUI_HOME_DIR="$HOME/webui"

# Clone the repository, the hard way (to allow existing linked folder of webui)
# Check if .git directory exists in WEBUI_HOME_DIR
if [ ! -d "${WEBUI_HOME_DIR}/.git" ]; then
    # Clone the repository, the hard way (to allow existing linked folder of webui)
    git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui /tmp/webui_temp
    mv /tmp/webui_temp/* /tmp/webui_temp/.git* "${WEBUI_HOME_DIR}"
    rm -rf /tmp/webui_temp
fi

# Check if it's the first time setup
if [ ! -f "ui-config.json" ]; then

    # Update and install perftools
    sudo apt-get update
    sudo apt-get install --no-install-recommends google-perftools -y

    cd "${WEBUI_HOME_DIR}"
    # If ui-config.json does not exist, it's the first time setup
    # Adding flags
    sed -i 's/#export COMMANDLINE_ARGS=""/export COMMANDLINE_ARGS="--share --no-half-vae --listen --api --xformers --no-download-sd-model --disable-nan-check"/' webui-user.sh
    wget https://huggingface.co/kiriyamaX/webui-configs/resolve/main/config_sdxl_v2.json -O config.json
    wget https://huggingface.co/kiriyamaX/webui-configs/resolve/main/ui-config_sdxl_v2.json -O ui-config.json

    # Install extensions
    cd extensions 
    git clone https://github.com/picobyte/stable-diffusion-webui-wd14-tagger
    git clone https://github.com/AlUlkesh/stable-diffusion-webui-images-browser
    git clone https://github.com/hako-mikan/sd-webui-supermerger
    git clone https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111
    git clone https://github.com/arenatemp/stable-diffusion-webui-model-toolkit  

    # ========== SECURITY: disallow gradio visiting models folder =============
    WEBUI_PY_FILE="${WEBUI_HOME_DIR}/webui.py"
    
    # Line to be added
    # https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cf2772fab0af5573da775e7437e6acdca424f26e/webui.py#L79C1-L96C10
    NEW_LINE="            blocked_paths=[\"/home/ubuntu/webui/models\"],"
    
    # Check if the line already exists to avoid duplicating it
    if ! grep -q "blocked_paths" "$WEBUI_PY_FILE"; then
        # Use sed to insert the line after the pattern
        sed -i "/app_kwargs={/i $NEW_LINE" "$WEBUI_PY_FILE"
    fi
fi

cd "${WEBUI_HOME_DIR}"
./webui.sh