#!/bin/sh
# tau-rag container entrypoint
# Dispatches on first arg: serve | migrate | init-corpus | shell | test
set -e

CMD="${1:-serve}"
shift || true

warmup_corpus() {
    if [ -n "$LAWDBHEB_ROOT" ] && [ -d "$LAWDBHEB_ROOT" ]; then
        echo "[entrypoint] Corpus root found: $LAWDBHEB_ROOT"
        # Optional: run init_corpus.py if no index exists yet.
        if [ ! -f "$TAU_RAG_RUNTIME_DIR/.corpus_loaded" ]; then
            echo "[entrypoint] First boot — loading corpus into index..."
            python -m scripts.init_corpus \
                --root "$LAWDBHEB_ROOT" \
                --subfolder "${LAWDBHEB_INDEX_SUBFOLDER:-legal_rag_index_hybrid}" \
                --runtime "$TAU_RAG_RUNTIME_DIR" \
            && touch "$TAU_RAG_RUNTIME_DIR/.corpus_loaded" \
            || echo "[entrypoint] WARN: corpus load failed — continuing with empty index"
        fi
    else
        echo "[entrypoint] LAWDBHEB_ROOT unset/missing — running with empty index"
    fi
}

seed_admin_key() {
    if [ -n "$TAU_RAG_SEED_ADMIN_KEY" ]; then
        echo "[entrypoint] Seeding admin API key from env"
        export TAU_RAG_SEEDED_ADMIN=1
    fi
}

case "$CMD" in
    serve)
        warmup_corpus
        seed_admin_key
        echo "[entrypoint] Starting tau-rag on ${TAU_RAG_HOST}:${TAU_RAG_PORT} "\
"(workers=${TAU_RAG_WORKERS}, preset=${TAU_RAG_PRESET})"
        exec gunicorn \
            --bind "${TAU_RAG_HOST}:${TAU_RAG_PORT}" \
            --workers "${TAU_RAG_WORKERS:-4}" \
            --worker-class uvicorn.workers.UvicornWorker \
            --timeout "${TAU_RAG_REQUEST_TIMEOUT_S:-30}" \
            --graceful-timeout 30 \
            --keep-alive 5 \
            --max-requests 1000 \
            --max-requests-jitter 50 \
            --access-logfile - \
            --error-logfile - \
            --log-level "${TAU_RAG_LOG_LEVEL:-info}" \
            tau_rag.api.fastapi_app:app
        ;;
    migrate)
        echo "[entrypoint] No schema migrations required for tau-rag (in-memory store)."
        exit 0
        ;;
    init-corpus)
        warmup_corpus
        ;;
    shell)
        exec python "$@"
        ;;
    test)
        export PYTHONPATH="/app${PYTHONPATH:+:$PYTHONPATH}"
        exec python -m pytest /app/tau_rag/tests "$@"
        ;;
    *)
        # Fall-through: run arbitrary command
        exec "$CMD" "$@"
        ;;
esac