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

if [ -d "/data" ]; then
  DEFAULT_STORAGE_ROOT="/data/rag_storage"
  DEFAULT_HF_HOME="/data/.huggingface"
else
  DEFAULT_STORAGE_ROOT="data/rag_storage"
  DEFAULT_HF_HOME="/tmp/.cache/huggingface"
fi

export HF_HOME="${HF_HOME:-$DEFAULT_HF_HOME}"
export LIGHTRAG_STORAGE_ROOT="${LIGHTRAG_STORAGE_ROOT:-$DEFAULT_STORAGE_ROOT}"

JURISDICTIONS="${JURISDICTIONS:-romania,bahrain}"
HOST="${LIGHTRAG_HOST:-127.0.0.1}"
ROOT="${LIGHTRAG_STORAGE_ROOT}"
GRAPHS="${LIGHTRAG_GRAPHS:-romania:9621,bahrain:9622}"
API_PORT="${PORT:-${API_PORT:-8000}}"

mkdir -p "${ROOT}"
mkdir -p "${HF_HOME}"

echo "📁 Storage root: ${ROOT}"
echo "💾 HF_HOME: ${HF_HOME}"
echo "🌍 JURISDICTIONS=${JURISDICTIONS}"
echo "🚀 LIGHTRAG_GRAPHS=${GRAPHS}"

echo "📥 Checking for knowledge graph data..."
missing=0
missing_jurs=()
IFS=',' read -r -a jurisdictions <<< "${JURISDICTIONS}"
for jur in "${jurisdictions[@]}"; do
  jur="$(echo "$jur" | xargs)"
  if [ ! -d "${ROOT}/${jur}" ]; then
    echo "   ❌ Missing: ${ROOT}/${jur}"
    missing=1
    missing_jurs+=("${jur}")
  else
    file_count=$(find "${ROOT}/${jur}" -type f | wc -l)
    echo "   ✅ Found: ${ROOT}/${jur} (${file_count} files)"
  fi
done

if [ "${missing}" -eq 1 ]; then
  echo ""
  echo "🚀 Downloading knowledge graph from Hugging Face..."
  echo "   🌍 Missing jurisdictions: ${missing_jurs[*]}"
  echo "   ⏳ Starting download process..."
  echo ""
  
  if python scripts/download_knowledge_graph.py; then
    echo ""
    echo "✅ Knowledge graph download complete"
  else
    echo ""
    echo "❌ Knowledge graph download failed"
    exit 1
  fi
else
  echo ""
  echo "✅ All knowledge graph data already present, skipping download"
fi

echo ""

PIDS=()
trap 'kill -TERM ${PIDS[@]:-} 2>/dev/null || true; wait ${PIDS[@]:-} 2>/dev/null || true' EXIT INT TERM

ENDPOINTS=()

IFS=',' read -r -a items <<< "${GRAPHS}"
for item in "${items[@]}"; do
  IFS=':' read -r id port <<< "${item}"
  dir="${ROOT}/${id}"
  mkdir -p "${dir}"

  echo "➡️  Start LightRAG '${id}' on ${HOST}:${port} (dir=${dir})"
  lightrag-server --host "${HOST}" --port "${port}" --working-dir "${dir}" &
  PIDS+=("$!")

  echo "   ⏳ Waiting health for ${id}..."
  ready=0
  for _ in {1..60}; do
    if curl -fsS "http://${HOST}:${port}/health" >/dev/null 2>&1; then
      echo "   ✅ ${id} ready"
      ready=1
      break
    fi
    sleep 2
  done

  if [ "${ready}" -ne 1 ]; then
    echo "   ❌ ${id} failed health check"
    exit 1
  fi

  ENDPOINTS+=("${id}=http://${HOST}:${port}")
done

export LIGHTRAG_ENDPOINTS="$(IFS=,; echo "${ENDPOINTS[*]}")"
export PORT="${API_PORT}"

echo "✅ LIGHTRAG_ENDPOINTS=${LIGHTRAG_ENDPOINTS}"
echo "🚀 Starting API on 0.0.0.0:${PORT} ..."
python agent_api.py