#!/bin/bash
set -e

# ═══════════════════════════════════════════════════════════════════════════════
# القيم تأتي من متغيرات البيئة المعرّفة في Dockerfile (أو من HF Spaces Secrets)
# يمكنك تجاوزها في واجهة HF Spaces تحت قسم "Repository secrets" دون تعديل الكود
# ═══════════════════════════════════════════════════════════════════════════════
MODELS_DIR="/data/models"
MODEL_PATH="$MODELS_DIR/$HF_FILE"

# ── التحقق من المتغيرات الأساسية ──────────────────────────────────────────────
if [ -z "$HF_REPO" ] || [ -z "$HF_FILE" ] || [ -z "$MODEL_NAME" ]; then
    echo "ERROR: HF_REPO, HF_FILE, and MODEL_NAME must all be set."
    exit 1
fi

echo "════════════════════════════════════════════"
echo "  Repo  : $HF_REPO"
echo "  File  : $HF_FILE"
echo "  Name  : $MODEL_NAME"
echo "════════════════════════════════════════════"

# إنشاء المجلدات المطلوبة
mkdir -p "$MODELS_DIR" /data/ollama /data/webui

# ── تحميل الملف من Hugging Face (يتخطى التحميل إن كان الملف موجوداً بالفعل) ──
if [ -f "$MODEL_PATH" ]; then
    echo "Model file already exists at $MODEL_PATH — skipping download."
else
    echo "Downloading $HF_FILE from $HF_REPO ..."

    # تمرير التوكن فقط إن كان مضبوطاً (للنماذج المحمية)
    HF_CLI_ARGS="--repo-type model --local-dir $MODELS_DIR"
    if [ -n "$HF_TOKEN" ]; then
        HF_CLI_ARGS="$HF_CLI_ARGS --token $HF_TOKEN"
    fi

    huggingface-cli download "$HF_REPO" "$HF_FILE" $HF_CLI_ARGS

    if [ ! -f "$MODEL_PATH" ]; then
        echo "ERROR: Download finished but file not found at $MODEL_PATH"
        exit 1
    fi
    echo "Download complete: $MODEL_PATH"
fi

# ── تشغيل خادم Ollama في الخلفية ──────────────────────────────────────────────
echo "Starting Ollama server..."
OLLAMA_MODELS=/data/ollama ollama serve &

# ── انتظار حتى يكون Ollama جاهزاً فعلاً (بدلاً من sleep ثابت) ────────────────
echo "Waiting for Ollama to be ready..."
MAX_WAIT=60
WAITED=0
until curl -s http://localhost:11434/api/tags > /dev/null 2>&1; do
    if [ "$WAITED" -ge "$MAX_WAIT" ]; then
        echo "ERROR: Ollama did not respond within ${MAX_WAIT}s."
        exit 1
    fi
    sleep 1
    WAITED=$((WAITED + 1))
done
echo "Ollama is ready (after ${WAITED}s)."

# ── إنشاء Modelfile يشير إلى الملف المحلي في المخزن ──────────────────────────
echo "Creating Modelfile..."
cat <<EOF > "$MODELS_DIR/Modelfile"
FROM $MODEL_PATH

PARAMETER num_ctx $MODEL_CTX
PARAMETER num_thread $MODEL_THREADS
EOF

# ── تسجيل النموذج في Ollama إن لم يكن مسجلاً بعد ───────────────────────────
if OLLAMA_MODELS=/data/ollama ollama list | grep -q "^$MODEL_NAME"; then
    echo "Model '$MODEL_NAME' already registered — skipping creation."
else
    echo "Registering model '$MODEL_NAME' in Ollama..."
    OLLAMA_MODELS=/data/ollama ollama create "$MODEL_NAME" -f "$MODELS_DIR/Modelfile"
    echo "Model registered successfully."
fi

# ── تحميل النموذج مسبقاً في الذاكرة (Warmup) ────────────────────────────────
# نرسل طلباً فارغاً بـ keep_alive: -1 فيُحمَّل النموذج الآن ويبقى في الذاكرة
# حتى إعادة التشغيل، فلا يوجد أي تأخير عند أول رسالة من المستخدم
echo "Pre-loading '$MODEL_NAME' into memory..."
curl -s -X POST http://localhost:11434/api/generate \
    -H "Content-Type: application/json" \
    -d "{\"model\": \"$MODEL_NAME\", \"prompt\": \"\", \"keep_alive\": -1}" \
    > /dev/null

if [ $? -eq 0 ]; then
    echo "Model '$MODEL_NAME' is loaded and ready in memory."
else
    echo "WARNING: Pre-load request failed. Model will load on first use instead."
fi

# ── تشغيل واجهة Open WebUI ────────────────────────────────────────────────────
echo "Starting Open WebUI..."
cd /app/backend
exec bash start.sh