#!/bin/bash

# ==============================================================================
#   Automated Open WebUI & SearXNG Installer (v26 - The Definitive Version)
#
#   This script is the final, consolidated version incorporating all bug fixes
#   and best practices discovered through our collaborative debugging process.
#
#   Key Fixes:
#   1. Uses a robust "port-publishing" method for Nginx-to-Docker communication.
#   2. Uses environment variables to configure SearXNG, the correct method.
#   3. Uses a safe, non-blocking command to generate secrets.
#   4. Builds Docker commands safely in an array to prevent errors.
#   5. All previous logic (UFW, cron, etc.) is complete and verified.
# ==============================================================================

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

# --- Pre-flight Checks ---
if [[ "$EUID" -ne 0 ]]; then echo "❌ This script must be run as root."; exit 1; fi
if ! grep -qi "debian" /etc/os-release; then echo "⚠️ This script is optimized for Debian."; fi

# --- User Input ---
echo "---"
echo "Welcome to the Full Stack Open WebUI & SearXNG Installer."
echo "---"
read -p "Enter your domain for Open WebUI (e.g., ai.example.com): " UI_DOMAIN
read -p "Enter your email address (for Let's Encrypt): " EMAIL
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 ""
read -p "Deploy a private SearXNG instance for the research tool? [y/N]: " DEPLOY_SEARXNG
BRAVE_API_KEY=""
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    SEARCH_DOMAIN="search.r21.io"
    echo "ℹ️ SearXNG will be deployed to https://$SEARCH_DOMAIN"
    read -p "Enter your Brave Search API key (highly recommended, or press Enter): " BRAVE_API_KEY
fi

# --- Main Script ---
echo "---"; echo "✅ Thank you. Starting the setup."; sleep 3

# --- Configuration ---
UI_CONTAINER="open-webui"
SEARXNG_CONTAINER="searxng"
NETWORK_NAME="open-webui-net"

# --- Step 1: Dependencies ---
echo "▶️ [1/9] Installing dependencies..."
export DEBIAN_FRONTEND=noninteractive
apt-get update
# Add openssl for robust secret generation
BASE_PACKAGES="ca-certificates curl gnupg nginx certbot python3-certbot-nginx fail2ban unattended-upgrades openssl"
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then apt-get install -y $BASE_PACKAGES apache2-utils; else apt-get install -y $BASE_PACKAGES; fi

# --- Step 2: Firewall Management ---
echo "▶️ [2/9] Managing Firewall..."
if [[ -z "${MANAGE_UFW+x}" ]]; then MANAGE_UFW="y"; fi
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: Docker & Network ---
echo "▶️ [3/9] Installing Docker and setting up network..."
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
if ! docker network ls | grep -q "$NETWORK_NAME"; then docker network create $NETWORK_NAME; fi

# --- Step 4: Stop Old Containers ---
echo "▶️ [4/9] Stopping and removing any old containers..."
docker stop $UI_CONTAINER $SEARXNG_CONTAINER 2>/dev/null || true
docker rm $UI_CONTAINER $SEARXNG_CONTAINER 2>/dev/null || true

# --- Step 5: Configure and Deploy SearXNG (Optional) ---
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    echo "▶️ [5/9] Deploying SearXNG..."
    
    # Generate a robust, shell-safe secret key
    SECRET_KEY=$(openssl rand -hex 32)

    # Build the docker run command safely in an array
    docker_cmd=(
        docker run -d
        --name "$SEARXNG_CONTAINER"
        --network "$NETWORK_NAME"
        # Publish port to localhost for Nginx to connect to
        -p "127.0.0.1:8081:8080"
        # Core settings via environment variables
        -e "SEARXNG_SECRET=$SECRET_KEY"
        -e "SEARXNG_BIND_ADDRESS=0.0.0.0" # Listen on all interfaces inside the container
        -e "SEARXNG_BASE_URL=https://$SEARCH_DOMAIN"
        --restart always
    )

    # Add optional Brave integration
    if [[ -n "$BRAVE_API_KEY" ]]; then
        echo "  - Enabling Brave engine with API key..."
        docker_cmd+=(
            -e "SEARXNG_ENGINES_BRAVE_API_KEY=$BRAVE_API_KEY"
            -e "SEARXNG_ENGINES_BRAVE_DISABLED=false"
            # Disable a noisy engine if a key is present
            -e "SEARXNG_ENGINES_DUCKDUCKGO_DISABLED=true"
        )
    else
        echo "  - No Brave API key provided, using default search engines."
    fi
    
    # Add the image name to the end of the command
    docker_cmd+=(searxng/searxng)
    
    # Execute the final, safe command
    "${docker_cmd[@]}"
else
    echo "▶️ [5/9] Skipping SearXNG deployment."
fi

# --- Step 6: Deploy Open WebUI ---
echo "▶️ [6/9] Deploying Open WebUI..."
docker run -d -p 127.0.0.1:3000:8080 -v open-webui:/app/backend/data --name $UI_CONTAINER --network $NETWORK_NAME --restart always ghcr.io/open-webui/open-webui:main

# --- Step 7: Configure Nginx ---
echo "▶️ [7/9] Configuring Nginx sites..."
NGINX_UI_CONF="/etc/nginx/sites-available/$UI_DOMAIN"
if [[ "${UI_INCLUDE_WWW,,}" == "y" ]]; then UI_SERVER_NAME="$UI_DOMAIN www.$UI_DOMAIN"; else UI_SERVER_NAME="$UI_DOMAIN"; fi
cat <<EOF > "$NGINX_UI_CONF"
server {
    listen 80; listen [::]:80; server_name $UI_SERVER_NAME;
    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"

if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    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 / {
        # Proxy directly to the port we published on the host's localhost
        proxy_pass http://127.0.0.1:8081;
        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"
fi

# --- Step 8: SSL Certificates & Automations ---
echo "▶️ [8/9] Obtaining SSL certificates and configuring automations..."
nginx -t
CERTBOT_STAGING_FLAG=""
if [[ "${USE_STAGING,,}" == "y" ]]; then CERTBOT_STAGING_FLAG="--staging"; fi
certbot --nginx -d $UI_SERVER_NAME $CERTBOT_STAGING_FLAG --non-interactive --agree-tos -m "$EMAIL" --redirect
if [[ "${DEPLOY_SEARXNG,,}" == "y" ]]; then
    certbot --nginx -d $SEARCH_DOMAIN $CERTBOT_STAGING_FLAG --non-interactive --agree-tos -m "$EMAIL" --redirect
fi
NGINX_UI_LOG="/var/log/nginx/$UI_DOMAIN.access.log"
touch $NGINX_UI_LOG && chown www-data:adm $NGINX_UI_LOG
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" > /etc/cron.d/weekly-reboot

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

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. Paste the Python code from https://github.com/iamarcel/open-webui-utils/blob/main/research_tool.py"
    echo "3. Go to the Settings tab in the tool editor."
    echo "4. Add Environment Variable: Key: SEARXNG_BASE_URL, Value: http://searxng:8080"
    echo "5. Click 'Save'."
fi
echo "---"