#!/bin/bash
set -e

# Path to the downloaded model
MODEL_DIR="/code/models"
MODEL_PATH="$MODEL_DIR/qwen2.5-3b-instruct-q4_k_m.gguf"

echo "=== ByteAstra Cloud Startup ==="

# 1. Download custom GGUF model from Hugging Face Hub using huggingface-cli
if [ ! -f "$MODEL_PATH" ]; then
    echo "Downloading official Qwen GGUF model from Hugging Face Hub..."
    mkdir -p "$MODEL_DIR"
    # Download the GGUF file using the modern hf CLI
    hf download Qwen/Qwen2.5-3B-Instruct-GGUF --include "qwen2.5-3b-instruct-q4_k_m.gguf" --local-dir "$MODEL_DIR"
    echo "Model downloaded successfully."
else
    echo "Model already exists at $MODEL_PATH"
fi

# 2. Start llama-cpp-python server in the background (runs on CPU, port 8001)
echo "Starting local llama-cpp-python server on port 8001..."
python -m llama_cpp.server \
    --model "$MODEL_PATH" \
    --port 8001 \
    --host 127.0.0.1 \
    --n_ctx 4096 \
    --n_threads 2 \
    --n_threads_batch 2 \
    --n_batch 256 \
    --chat_format chatml &

# 3. Wait for the local llama.cpp server to be ready
echo "Waiting for llama.cpp server to initialize..."
until curl -s http://127.0.0.1:8001/v1/models > /dev/null; do
    sleep 2
    echo "Still waiting..."
done
echo "llama.cpp server is ready!"

# 4. Set environment variables to point FastAPI to the local llama.cpp server
export OPENAI_BASE_URL="http://127.0.0.1:8001/v1"
export LLM_BASE_URL="http://127.0.0.1:8001/v1"
export MODEL_NAME="qwen2.5-3b-instruct-q4_k_m.gguf"
export LLM_MODEL_NAME="qwen2.5-3b-instruct-q4_k_m.gguf"
export LLM_MAX_TOKENS=2048


# 4.5. Extract pre-built ChromaDB index from zip if present
if [ -f "chroma_store.zip" ]; then
    echo "Extracting pre-built ChromaDB index from zip..."
    python -c "import zipfile; zipfile.ZipFile('chroma_store.zip').extractall('chroma_store')"
    echo "Extraction complete."
fi

# 5. Run ingest only if the collection is empty or new files exist
echo "Checking knowledge base..."
CHROMA_COUNT=$(python -c "
import sys
sys.path.insert(0, '.')
try:
    from app.services.rag import get_or_create_collection
    col = get_or_create_collection('domain_ayurveda')
    print(col.count())
except Exception as e:
    print(0)
" 2>/dev/null || echo 0)
echo "ChromaDB has $CHROMA_COUNT chunks indexed."
if [ "$CHROMA_COUNT" -lt "100" ]; then
    echo "Indexing knowledge base (first run or collection is empty)..."
    python scripts/ingest.py --domain ayurveda --source-dir ./data/ayurveda || echo "Ingest step done."
else
    echo "Knowledge base already indexed ($CHROMA_COUNT chunks). Skipping ingest."
fi

# 6. Start FastAPI server on port 7860 (Hugging Face default)
echo "Starting FastAPI backend server on port 7860..."
exec uvicorn app.main:app --host 0.0.0.0 --port 7860