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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
SANDBOX_DIR="${ROOT_DIR}/sandbox"
VENV_DIR="${SANDBOX_DIR}/benchmark/.venv"

# shellcheck source=common/docker.sh
source "${SANDBOX_DIR}/common/docker.sh"
# shellcheck source=common/python.sh
source "${SANDBOX_DIR}/common/python.sh"

DATASET="${1:-all}"
PG_PORT="${PGGRAPH_PG_PORT:-55432}"
CONTAINER_NAME="${PGGRAPH_CONTAINER_NAME:-pggraph-sandbox}"
IMAGE_NAME="${PGGRAPH_IMAGE_NAME:-pggraph-postgres:17}"
shift || true

if [ "${DATASET}" = "cleanup" ]; then
  CLEANUP_TARGETS=()
  CLEANUP_FLAGS=()
  for arg in "$@"; do
    case "${arg}" in
      datasets|results|venv|docker|all) CLEANUP_TARGETS+=("${arg}") ;;
      *) CLEANUP_FLAGS+=("${arg}") ;;
    esac
  done
  if [ "${#CLEANUP_TARGETS[@]}" -eq 0 ]; then
    CLEANUP_TARGETS=(all)
  fi
  python3 "${SANDBOX_DIR}/common/run_benchmarks.py" \
    --cleanup "${CLEANUP_TARGETS[@]}" \
    --container "${CONTAINER_NAME}" \
    --image "${IMAGE_NAME}" \
    --datasets-dir "${SANDBOX_DIR}/benchmark/datasets" \
    --results-dir "${SANDBOX_DIR}/benchmark/results" \
    --benchmark-venv "${SANDBOX_DIR}/benchmark/.venv" \
    --playground-venv "${SANDBOX_DIR}/playground/.venv" \
    ${CLEANUP_FLAGS[@]+"${CLEANUP_FLAGS[@]}"}
  exit 0
fi

require_docker
ensure_pggraph_image "${ROOT_DIR}" "${IMAGE_NAME}"
ensure_pggraph_container "${CONTAINER_NAME}" "${IMAGE_NAME}" "${PG_PORT}"

BENCHMARK_PYTHON=""
for candidate in python3.13 python3.12 python3.11 python3.10 python3; do
  if command -v "${candidate}" >/dev/null 2>&1; then
    candidate_path="$(command -v "${candidate}")"
    if "${candidate_path}" -c 'import sys; raise SystemExit(0 if sys.version_info >= (3, 10) else 1)' >/dev/null 2>&1; then
      BENCHMARK_PYTHON="$(pggraph_resolve_python "${candidate_path}")"
      break
    fi
  fi
done

if [ -z "${BENCHMARK_PYTHON}" ]; then
  echo "Error: benchmark timing requires Python 3.10 or newer." >&2
  exit 1
fi

if [ -d "${VENV_DIR}" ]; then
  if ! "${VENV_DIR}/bin/python" -c 'import sys; raise SystemExit(0 if sys.version_info >= (3, 10) else 1)' >/dev/null 2>&1; then
    rm -rf "${VENV_DIR}"
  fi
fi

if [ ! -d "${VENV_DIR}" ]; then
  "${BENCHMARK_PYTHON}" -m venv "${VENV_DIR}"
fi

pggraph_prepare_venv_requirements \
  "${VENV_DIR}" \
  "${SANDBOX_DIR}/benchmark/requirements.txt"

"${VENV_DIR}/bin/python" "${SANDBOX_DIR}/common/run_benchmarks.py" \
  --dataset "${DATASET}" \
  --container "${CONTAINER_NAME}" \
  --image "${IMAGE_NAME}" \
  --host 127.0.0.1 \
  --port "${PG_PORT}" \
  --database postgres \
  --user postgres \
  --password postgres \
  --datasets-dir "${SANDBOX_DIR}/benchmark/datasets" \
  --results-dir "${SANDBOX_DIR}/benchmark/results" \
  --benchmark-venv "${SANDBOX_DIR}/benchmark/.venv" \
  --playground-venv "${SANDBOX_DIR}/playground/.venv" \
  "$@"
