#!/usr/bin/env bash
set -euo pipefail

# run.sh — friendly wrapper for running the distillation pipeline

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$ROOT"

# -----------------------------
# Default config (can be overridden in env or CLI)
# -----------------------------
ES_NODE="${ES_NODE:-http://localhost:9200}"
ES_INDEX="${ES_INDEX:-quo_distill_index}"

PIPELINE_SEED_MODE="${PIPELINE_SEED_MODE:-question-first}"

RETRIEVAL_MODE="${RETRIEVAL_MODE:-hybrid}"
RETRIEVAL_K="${RETRIEVAL_K:-6}"

QUESTION_MAX="${QUESTION_MAX:-3}"
LIMIT="${LIMIT:-5}"

GENERATOR_MODEL="${GENERATOR_MODEL:-qwen3-vl:8b-thinking}"
VERIFIER_MODEL="${VERIFIER_MODEL:-tensortemplar/patronus-lynx:8b-instruct-q4_K_M}"
REWARD_MODEL="${REWARD_MODEL:-tensortemplar/patronus-lynx:8b-instruct-q4_K_M}"

VERBOSE=0
THINKING=0

# -----------------------------
# Argument parsing
# -----------------------------
while [[ $# -gt 0 ]]; do
  case "$1" in
    -l|--limit)
      LIMIT="$2"
      shift 2
      ;;
    -q|--questions)
      QUESTION_MAX="$2"
      shift 2
      ;;
    -v|--verbose)
      VERBOSE=1
      shift
      ;;
    -t|--thinking)
      THINKING=1
      shift
      ;;
    *)
      echo "Unknown option: $1"
      echo "Usage: ./run.sh [--limit N] [--questions N] [--verbose] [--thinking]"
      exit 1
      ;;
  esac
done

# If thinking flag is present, use thinking model + thinking prompt header
if [[ "$THINKING" -eq 1 ]]; then
  echo "🧠 Using THINKING MODE"
  GENERATOR_MODEL="qwen3-vl:8b-thinking"
fi

# -----------------------------
# Export for pipeline
# -----------------------------
export ES_NODE
export ES_INDEX
export PIPELINE_SEED_MODE

export RETRIEVAL_MODE
export RETRIEVAL_K

export QUESTION_MAX
export LIMIT

export GENERATOR_PROVIDER="ollama"
export GENERATOR_MODEL

export VERIFIER_PROVIDER="ollama"
export VERIFIER_MODEL

export REWARD_PROVIDER="ollama"
export REWARD_MODEL

echo "========================================"
echo "🚀 Running Distillation Pipeline"
echo " ES_NODE:          $ES_NODE"
echo " ES_INDEX:         $ES_INDEX"
echo " Retrieval:        $RETRIEVAL_MODE (k=$RETRIEVAL_K)"
echo " Mode:             $PIPELINE_SEED_MODE"
echo " Limit:            $LIMIT"
echo " Question max:     $QUESTION_MAX"
echo " Generator model:  $GENERATOR_MODEL"
echo " Verifier model:   $VERIFIER_MODEL"
echo " Reward model:     $REWARD_MODEL"
echo " Verbose:          $VERBOSE"
echo "========================================"
echo

if [[ "$VERBOSE" -eq 1 ]]; then
  npm run pipeline -- --limit "$LIMIT" --verbose
else
  npm run pipeline -- --limit "$LIMIT"
fi