#!/bin/bash
# QAFD-RAG Benchmark Runner
# Unified interface for all benchmarks

set -euo pipefail

cd "$(dirname "${BASH_SOURCE[0]}")" || { echo "Error: Cannot cd to script directory"; exit 1; }

# Conda environment (has torch, transformers, sentence-transformers, etc.)
CONDA_ENV="${QAFD_CONDA_ENV:-qafd-rag}"
PYTHON="conda run --live-stream -n ${CONDA_ENV} python"

# Set NVIDIA library paths (needed for torch with pip-installed CUDA libs)
NVIDIA_LIB="$(conda info --base)/envs/${CONDA_ENV}/lib/python3.10/site-packages/nvidia"
if [ -d "$NVIDIA_LIB" ]; then
    export LD_LIBRARY_PATH="${NVIDIA_LIB}/cudnn/lib:${NVIDIA_LIB}/cublas/lib:${NVIDIA_LIB}/cuda_runtime/lib:${NVIDIA_LIB}/cuda_cupti/lib:${NVIDIA_LIB}/cuda_nvrtc/lib:${NVIDIA_LIB}/cufft/lib:${NVIDIA_LIB}/curand/lib:${NVIDIA_LIB}/cusolver/lib:${NVIDIA_LIB}/cusparse/lib:${NVIDIA_LIB}/nccl/lib:${NVIDIA_LIB}/nvjitlink/lib:${LD_LIBRARY_PATH:-}"
fi

# Load API key from .env if not already set
if [ -z "${OPENAI_API_KEY:-}" ]; then
    if [ -f .env ]; then
        eval "$(grep -E '^OPENAI_API_KEY=' .env | head -1)"
        export OPENAI_API_KEY
    fi
    if [ -z "${OPENAI_API_KEY:-}" ]; then
        echo "Error: OPENAI_API_KEY not set. Export it or add it to a .env file."
        exit 1
    fi
fi

# Show help
show_help() {
    echo "Usage: ./run.sh <task> [options]"
    echo ""
    echo "Tasks (graph type):"
    echo "  ultradomain      QA benchmark on UltraDomain dataset       [entity graph]"
    echo "  text2sql         Text-to-SQL benchmark                     [entity graph]"
    echo "  summarization    Document summarization QA benchmark        [entity graph]"
    echo "  multihop         Multi-hop QA benchmark (MuSiQue, etc.)    [entity graph]"
    echo ""
    echo "Note: run.sh uses the entity graph pipeline. For the passage-entity graph"
    echo "  (recommended for multihop), use the unified runner instead:"
    echo "  python benchmarks/run.py --task multihop --dataset musique --questions 100"
    echo ""
    echo "Common Options (all tasks):"
    echo "  --questions N        Number of questions to benchmark"
    echo "  --build              Build KG only, don't run benchmark"
    echo "  --force-build        Force rebuild KG even if exists"
    echo "  --max-documents N    Max documents for KG building"
    echo "  --embedding MODEL    Embedding model (openai-small, openai-large, jina-v3)"
    echo "  --llm MODEL          LLM model (gpt-4o-mini, gpt-4o, gpt-5-nano, gpt-5-mini, gpt-5, gpt-oss-120b)"
    echo ""
    echo "Summarization-specific Options:"
    echo "  --dataset NAME       Dataset: squality (default: squality)"
    echo ""
    echo "Multihop-specific Options:"
    echo "  --dataset NAME       Dataset: musique, hotpotqa, 2wikimultihopqa"
    echo ""
    echo "Text2SQL-specific Options:"
    echo "  --benchmark NAME     Benchmark: spider2-lite (default) or bird"
    echo "  --db NAME            Database name (e.g., Pagila, superhero)"
    echo ""
    echo "Text2SQL End-to-End Pipeline:"
    echo "  Place your .sqlite file in data/text2sql/<benchmark>/databases/<DB>/<DB>.sqlite"
    echo "  The DB summary (JSON) is auto-generated on first run."
    echo ""
    echo "Examples:"
    echo "  # Build KG"
    echo "  ./run.sh ultradomain --build --max-documents 100"
    echo "  ./run.sh summarization --build --max-documents 10"
    echo "  ./run.sh multihop --dataset musique --build --max-documents 500"
    echo "  ./run.sh multihop --dataset hotpotqa --build --max-documents 500"
    echo ""
    echo "  # Run benchmark"
    echo "  ./run.sh ultradomain --questions 10"
    echo "  ./run.sh summarization --questions 50"
    echo "  ./run.sh multihop --dataset musique --questions 100"
    echo "  ./run.sh multihop --dataset hotpotqa --questions 100"
    echo "  ./run.sh multihop --dataset 2wikimultihopqa --questions 100"
    echo "  ./run.sh text2sql --questions 5 --db Pagila"
    echo "  ./run.sh text2sql --benchmark bird --questions 5 --db superhero"
    echo ""
    echo "  # Use a local OpenAI-compatible model"
    echo "  export LOCAL_LLM_BASE_URL='http://localhost:8000/v1'"
    echo "  ./run.sh ultradomain --questions 10 --llm gpt-oss-120b"
    echo "  ./run.sh multihop --dataset musique --questions 10 --llm gpt-oss-120b"
}

case "${1:-help}" in
    ultradomain)
        shift
        $PYTHON benchmarks/ultradomain/benchmark_ultradomain.py "$@"
        ;;
    text2sql)
        shift
        $PYTHON benchmarks/text2sql/benchmark_text2sql.py "$@"
        ;;
    summarization)
        shift
        $PYTHON benchmarks/summarization/benchmark_summarization.py "$@"
        ;;
    multihop)
        shift
        $PYTHON benchmarks/multihop/benchmark_multihop.py "$@"
        ;;
    help|--help|-h)
        show_help
        ;;
    *)
        echo "Error: Unknown task '${1}'. Run './run.sh --help' for usage."
        exit 1
        ;;
esac