#!/bin/bash
set -eo pipefail

echo "=========================================================="
echo " Starting Ollama + Nginx on Hugging Face Space (T4-small) "
echo "=========================================================="

# 1. GPU & CUDA Environment Validation
echo "[1/5] Validating GPU and CUDA environment..."
if ! python3 /scripts/gpu_detect.py; then
    echo "[CRITICAL ERROR] GPU environment validation failed."
    echo "Deployment requires adequate GPU (>= ${MIN_VRAM_MB:-15000} MB VRAM) and CUDA >= ${MIN_CUDA_MAJOR:-11}.${MIN_CUDA_MINOR:-8}."
    echo "Aborting startup to prevent unaccelerated execution."
    exit 1
fi
echo "[OK] GPU environment verified successfully."

# 2. Authentication Configuration
echo "[2/5] Configuring Nginx Basic Authentication..."
if [ -z "$NGINX_PASSWORD" ]; then
    echo "[CRITICAL ERROR] NGINX_PASSWORD environment variable is not defined."
    echo "Please configure NGINX_PASSWORD in Hugging Face Space Secrets."
    exit 1
fi

python3 /scripts/generate_htpasswd.py
echo "[OK] Nginx authentication configured."

# 3. Start Ollama daemon
echo "[3/5] Starting Ollama daemon..."
export OLLAMA_HOST="${OLLAMA_HOST:-0.0.0.0}"
export OLLAMA_PORT="${OLLAMA_PORT:-11434}"
export OLLAMA_ORIGINS="${OLLAMA_ORIGINS:-*}"
export OLLAMA_MODELS="${OLLAMA_MODELS:-/data/ollama}"

# Configure and verify persistent model storage
if mkdir -p "$OLLAMA_MODELS" 2>/dev/null && [ -w "$OLLAMA_MODELS" ]; then
    echo "[OK] Persistent model storage active at ${OLLAMA_MODELS}."
else
    echo "[WARNING] Cannot write to ${OLLAMA_MODELS}, falling back to /home/ollama/.ollama/models"
    export OLLAMA_MODELS="/home/ollama/.ollama/models"
    mkdir -p "$OLLAMA_MODELS" 2>/dev/null || true
fi

ollama serve &
OLLAMA_PID=$!

# Wait for Ollama to become healthy
echo "Waiting for Ollama to become ready..."
MAX_RETRIES=30
COUNTER=0
until curl -s "http://127.0.0.1:${OLLAMA_PORT}/api/version" >/dev/null 2>&1; do
    sleep 1
    COUNTER=$((COUNTER + 1))
    if [ $COUNTER -ge $MAX_RETRIES ]; then
        echo "[CRITICAL ERROR] Ollama failed to start within ${MAX_RETRIES} seconds."
        kill -9 $OLLAMA_PID 2>/dev/null || true
        exit 1
    fi
done
echo "[OK] Ollama is up and running."

# 4. Pull required model
TARGET_MODEL="${MODEL_NAME:-qwen2.5:14b}"
echo "[4/5] Pulling Ollama model: ${TARGET_MODEL}..."
ollama pull "${TARGET_MODEL}"
echo "[OK] Model ${TARGET_MODEL} is ready."

# 5. Trap signals for graceful shutdown
cleanup() {
    echo "Received termination signal. Shutting down gracefully..."
    nginx -s quit 2>/dev/null || true
    kill -TERM "$OLLAMA_PID" 2>/dev/null || true
    wait "$OLLAMA_PID" 2>/dev/null || true
    echo "Services stopped."
    exit 0
}
trap cleanup SIGTERM SIGINT

# 6. Start Nginx reverse proxy
echo "[5/5] Starting Nginx reverse proxy on port ${NGINX_PORT:-7860}..."
mkdir -p /tmp/nginx/body /tmp/nginx/proxy /tmp/nginx/fastcgi /tmp/nginx/uwsgi /tmp/nginx/scgi /var/log/nginx
nginx -g "daemon off;" &
NGINX_PID=$!

# Wait on foreground processes
wait -n $NGINX_PID $OLLAMA_PID