#!/bin/bash

# ==============================================================================
#   Automated Open WebUI & SearXNG Installer (v12 - GPG-Powered)
#
#   This script will:
#   1. Use `gpg` for random data generation, removing the `openssl` dependency.
#   2. Correctly configure Nginx to resolve Docker container names.
#   3. Deploy a complete, secure, and automated stack for Open WebUI and SearXNG.
# ==============================================================================

# --- Safety Checks ---
set -euo pipefail

# --- Pre-flight Checks ---
if [[ "$EUID" -ne 0 ]]; then
  echo "❌ This script must be run as root. Please use 'sudo ./onepush.sh'"
  exit 1
fi
if ! grep -qi "debian" /etc/os-release; then
    echo "⚠️ This script is optimized for Debian. Running on another OS may have unintended results."
fi

# --- User Input ---
echo "---"
echo "Welcome to the Full Stack Open WebUI & SearXNG Installer."
echo "---"
echo "--- Part 1: Open WebUI Configuration ---"
read -p "Enter your domain for Open WebUI (e.g., ai.example.com): " UI_DOMAIN
if [[ -z "$UI_DOMAIN" ]]; then echo "❌ Domain name cannot be empty."; exit 1; fi

read -p "Enter your email address (for Let's Encrypt): " EMAIL
if [[ -z "$EMAIL" ]]; then echo "❌ Email address cannot be empty."; exit 1; fi

read -p "Also secure the 'www' version of this domain (www.$UI_DOMAIN)? [y/N]: " UI_INCLUDE_WWW
read -p "Use Let's Encrypt staging environment (for testing)? [y/N]: " USE_STAGING

echo ""
if command -v ufw &> /dev/null; then
    echo "ℹ️ UFW firewall is detected on this system."
    read -p "Allow this script to manage UFW rules (allow HTTP/HTTPS)? [Y/n]: " MANAGE_UFW
else
    read -p "Install and configure UFW firewall? (Recommended) [Y/n]: " MANAGE_UFW
fi

echo ""
echo "--- Part 2: Optional SearXNG Search Engine ---"
read -p "Deploy a private SearXNG instance for the research tool? [y/N]: " DEPLOY_SEARXNG

BRAVE_API_KEY=""
RESEARCH_TOOL_URL=""
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    SEARCH_DOMAIN="search.r21.io"
    echo "ℹ️ SearXNG will be deployed to https://$SEARCH_DOMAIN"
    echo "You can provide a Brave Search API key for enhanced, ad-free results."
    read -p "Enter your Brave Search API key (or press Enter to skip): " BRAVE_API_KEY

    if [[ -z "$BRAVE_API_KEY" ]]; then
        echo ""
        echo "You can provide a URL to a custom research_tool.py file."
        read -p "Enter URL for a custom research_tool.py (or press Enter for default): " RESEARCH_TOOL_URL
        if [[ -n "$RESEARCH_TOOL_URL" ]]; then
            echo "Validating URL..."
            if ! curl --head --fail -sL "$RESEARCH_TOOL_URL" > /dev/null || [[ ! "$RESEARCH_TOOL_URL" == *.py ]]; then
                echo "❌ Invalid or unreachable URL, or it doesn't end in .py. Proceeding with default."
                RESEARCH_TOOL_URL=""
            else
                echo "✅ URL is valid."
            fi
        fi
    fi
fi

echo "---"
echo "✅ Thank you. Starting the setup."
sleep 3

# --- Configuration Variables ---
UI_CONTAINER="open-webui"
UI_VOLUME="open-webui"
NGINX_UI_CONF="/etc/nginx/sites-available/$UI_DOMAIN"
NGINX_UI_LOG="/var/log/nginx/$UI_DOMAIN.access.log"

# --- Step 1: System Preparation and Dependencies ---
echo "▶️ [1/11] Updating system and installing dependencies..."
export DEBIAN_FRONTEND=noninteractive
apt-get update
# MODIFIED: Removed 'openssl' from the dependency list. 'gnupg' is already installed.
BASE_PACKAGES="ca-certificates curl gnupg nginx certbot python3-certbot-nginx fail2ban unattended-upgrades"
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    apt-get install -y $BASE_PACKAGES apache2-utils
else
    apt-get install -y $BASE_PACKAGES
fi

# --- Step 2: Configure Firewall ---
echo "▶️ [2/11] Managing Firewall..."
if [[ "${MANAGE_UFW,,}" != "n" ]]; then
    if ! command -v ufw &> /dev/null; then apt-get install -y ufw; fi
    ufw allow ssh > /dev/null; ufw allow 'Nginx Full' > /dev/null; ufw --force enable
else
    if command -v ufw &> /dev/null && ufw status | grep -q "Status: active"; then
        echo "  - ⚠️ UFW is active but not managed. Disabling to prevent setup failure."
        ufw disable
    fi
fi

# --- Step 3: Install Docker ---
echo "▶️ [3/11] Installing Docker..."
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/debian/gpg | gpg --dearmor -o /etc/apt/keyrings/docker.gpg
chmod a+r /etc/apt/keyrings/docker.gpg
echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/debian $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | tee /etc/apt/sources.list.d/docker.list > /dev/null
apt-get update
apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

# --- Step 4: Create Shared Docker Network ---
echo "▶️ [4/11] Setting up Docker network..."
if ! docker network ls | grep -q "open-webui-net"; then
    docker network create open-webui-net
fi

# --- Step 5: Deploy Open WebUI Container ---
echo "▶️ [5/11] Deploying Open WebUI container..."
docker stop $UI_CONTAINER 2>/dev/null || true
docker rm $UI_CONTAINER 2>/dev/null || true
docker run -d -p 127.0.0.1:3000:8080 \
  --add-host=host.docker.internal:host-gateway \
  -v $UI_VOLUME:/app/backend/data \
  --name $UI_CONTAINER \
  --network open-webui-net \
  --restart always \
  ghcr.io/open-webui/open-webui:main

# --- Step 6: Configure Nginx for Open WebUI ---
echo "▶️ [6/11] Configuring Nginx for Open WebUI..."
if [[ "${UI_INCLUDE_WWW,,}" == "y" ]]; then NGINX_UI_SERVER_NAME="$UI_DOMAIN www.$UI_DOMAIN"; else NGINX_UI_SERVER_NAME="$UI_DOMAIN"; fi
cat <<EOF > "$NGINX_UI_CONF"
server {
    listen 80; listen [::]:80;
    server_name $NGINX_UI_SERVER_NAME;
    access_log $NGINX_UI_LOG; error_log /var/log/nginx/$UI_DOMAIN.error.log;
    location / {
        proxy_pass http://127.0.0.1:3000;
        proxy_set_header Host \$host; proxy_set_header X-Real-IP \$remote_addr;
        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for; proxy_set_header X-Forwarded-Proto \$scheme;
        proxy_http_version 1.1; proxy_set_header Upgrade \$http_upgrade; proxy_set_header Connection "upgrade";
    }
}
EOF
ln -sfn "$NGINX_UI_CONF" "/etc/nginx/sites-enabled/$UI_DOMAIN"

# --- Step 7: Deploy SearXNG (Optional) ---
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    echo "▶️ [7/11] Deploying SearXNG..."
    
    sudo mkdir -p /srv/searxng
    
    SECRET_KEY=$(gpg --gen-random --armor 1 24)
    
    # Dynamically create settings.yml
    {
    echo "server:"
    echo "  https: false"
    echo "  secret_key: \"$SECRET_KEY\""
    echo ""
    echo "search:"
    echo "  default_lang: en"
    echo ""
    echo "engines:"
    if [[ -n "$BRAVE_API_KEY" ]]; then
        # MODIFIED: Changed engine name from "brave_search" to "brave"
        echo "  - name: brave"
        echo "    engine: brave"
        echo "    api_key: \"$BRAVE_API_KEY\""
        echo "    shortcut: b"
        echo "    disabled: false"
    fi
    echo "  - name: google"
    echo "    engine: google"
    echo "    disabled: false"
    } | sudo tee /srv/searxng/settings.yml > /dev/null

    SEARXNG_CONTAINER="searxng"
    docker stop $SEARXNG_CONTAINER 2>/dev/null || true
    docker rm $SEARXNG_CONTAINER 2>/dev/null || true
    docker run -d \
      --name $SEARXNG_CONTAINER \
      --network open-webui-net \
      -v /srv/searxng:/etc/searxng \
      --restart always \
      searxng/searxng
else
    echo "▶️ [7/11] Skipping SearXNG deployment."
fi

# --- Step 8: Configure Nginx for SearXNG (Optional) ---
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    echo "▶️ [8/11] Configuring Nginx for SearXNG..."
    NGINX_SEARCH_CONF="/etc/nginx/sites-available/$SEARCH_DOMAIN"
    
    echo "Please set a password for the public SearXNG instance."
    sudo htpasswd -c /etc/nginx/.htpasswd admin
    
    cat <<EOF > "$NGINX_SEARCH_CONF"
server {
    listen 80; listen [::]:80;
    server_name $SEARCH_DOMAIN;
    location / {
        resolver 127.0.0.11;
        set \$searxng_upstream http://searxng:8080;
        proxy_pass \$searxng_upstream;
        auth_basic "Private Search Instance";
        auth_basic_user_file /etc/nginx/.htpasswd;
        proxy_set_header Host \$host;
        proxy_set_header X-Real-IP \$remote_addr;
        proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for;
        proxy_set_header X-Forwarded-Proto \$scheme;
        proxy_set_header X-Forwarded-Host \$server_name;
    }
}
EOF
    ln -sfn "$NGINX_SEARCH_CONF" "/etc/nginx/sites-enabled/$SEARCH_DOMAIN"
else
    echo "▶️ [8/11] Skipping Nginx config for SearXNG."
fi

# --- Step 9: Obtain SSL Certificates ---
echo "▶️ [9/11] Finalizing Nginx and obtaining SSL certificates..."
nginx -t
CERTBOT_STAGING_FLAG=""
if [[ "${USE_STAGING,,}" == "y" ]]; then CERTBOT_STAGING_FLAG="--staging"; fi

# Get cert for UI Domain
if [[ "${UI_INCLUDE_WWW,,}" == "y" ]]; then UI_DOMAINS="-d $UI_DOMAIN -d www.$UI_DOMAIN"; else UI_DOMAINS="-d $UI_DOMAIN"; fi
certbot --nginx $UI_DOMAINS $CERTBOT_STAGING_FLAG --non-interactive --agree-tos -m "$EMAIL" --redirect

# Get cert for Search Domain if deployed
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    certbot --nginx -d $SEARCH_DOMAIN $CERTBOT_STAGING_FLAG --non-interactive --agree-tos -m "$EMAIL" --redirect
fi

# --- Step 10: Configure Automations ---
echo "▶️ [10/11] Configuring system and service automations..."
cat <<EOF > /etc/fail2ban/filter.d/open-webui.conf
[Definition]
failregex = ^<HOST> .* "POST /api/v1/auths/signin.*" 400
ignoreregex =
EOF
cat <<EOF > /etc/fail2ban/jail.d/open-webui.local
[open-webui]
enabled = true; port = http,https; filter = open-webui;
logpath = $NGINX_UI_LOG; maxretry = 5; findtime = 3600; bantime = 86400;
EOF
systemctl restart fail2ban
dpkg-reconfigure -plow unattended-upgrades
if ! [ "$(docker ps -q -f name=watchtower)" ]; then
    docker run -d --name watchtower -v /var/run/docker.sock:/var/run/docker.sock --restart always containrrr/watchtower
fi
echo "0 9 * * 5 root /sbin/reboot # Weekly reboot for system health" > /etc/cron.d/weekly-reboot

# --- Step 11: Finalization ---
echo "▶️ [11/11] Finalizing..."
systemctl restart nginx
systemctl restart fail2ban

echo "---"
echo "✅🎉 **DEPLOYMENT COMPLETE!** 🎉✅"
echo ""
echo "--- ACCESS ---"
echo "  - Open WebUI:  https://$UI_DOMAIN"
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
echo "  - SearXNG:     https://$SEARCH_DOMAIN (user: admin)"
fi
echo ""
echo "--- NEXT STEPS: ADDING THE RESEARCH TOOL ---"
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    echo "1. Go to Open WebUI -> Settings -> Tools."
    echo "2. Click 'Add Tool' and paste the Python code below into the editor."
    echo "3. Open the 'Settings' tab in the tool editor."
    echo "4. In the 'Environment Variables' section, add:"
    echo "   - Key:   SEARXNG_BASE_URL"
    echo "   - Value: http://searxng:8080"
    echo "5. Click 'Save' to finish."
    echo ""
    echo "--- Python Code for Research Tool ---"
    if [[ -n "$RESEARCH_TOOL_URL" ]]; then
        echo "Paste the content from your custom URL: $RESEARCH_TOOL_URL"
    else
        echo "Paste the default code from here: https://github.com/iamarcel/open-webui-utils/blob/main/research_tool.py"
    fi
else
    echo "SearXNG was not deployed, so no research tool was configured."
fi
echo "---"