#!/usr/bin/env bash
# One-time venv for edge Jetson counter — NVIDIA torch required (not PyPI).
set -euo pipefail

INSTALL_DIR="${INSTALL_DIR:-/opt/jetson-counter}"
VENV_DIR="${VENV_DIR:-/opt/jetson-counter/venv}"
SERVICE_USER="${SERVICE_USER:-jetson}"
TORCH_WHEEL_URL="${TORCH_WHEEL_URL:-https://developer.download.nvidia.com/compute/redist/jp/v60/pytorch/torch-2.4.0a0+3bcc3cddb5.nv24.07.16234504-cp310-cp310-linux_aarch64.whl}"

if [[ "$(id -u)" -ne 0 ]]; then
  echo "Run as root: sudo ./setup-venv.sh"
  exit 1
fi

apt-get install -y libopenblas-base libopenmpi-dev libomp-dev 2>/dev/null || true

mkdir -p "${INSTALL_DIR}"
if [[ ! -x "${VENV_DIR}/bin/python" ]]; then
  python3 -m venv --system-site-packages "${VENV_DIR}"
fi
chown -R "${SERVICE_USER}:${SERVICE_USER}" "${INSTALL_DIR}"

sudo -u "${SERVICE_USER}" "${VENV_DIR}/bin/pip" install --upgrade pip
sudo -u "${SERVICE_USER}" "${VENV_DIR}/bin/pip" install "numpy<2"

if ! sudo -u "${SERVICE_USER}" PYTHONNOUSERSITE=1 "${VENV_DIR}/bin/python" -c "import torch; assert torch.cuda.is_available()" 2>/dev/null; then
  echo "Installing NVIDIA Jetson torch wheel..."
  sudo -u "${SERVICE_USER}" "${VENV_DIR}/bin/pip" install --no-cache-dir "${TORCH_WHEEL_URL}"
fi

sudo -u "${SERVICE_USER}" "${VENV_DIR}/bin/pip" install ultralytics flask opencv-python

sudo -u "${SERVICE_USER}" PYTHONNOUSERSITE=1 "${VENV_DIR}/bin/python" -c "
import torch
from ultralytics import YOLO
import cv2
import flask
print('venv ok | torch', torch.__version__, '| cuda', torch.cuda.is_available())
"

echo ""
echo "If torchvision import fails for ultralytics, copy/build torchvision into venv."
echo "Next: cp config.env.example .env && nano .env && sudo ./install-services.sh"