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

PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
VENV_PATH="${PROJECT_ROOT}/.venv-train"
PYTHON_BIN="${PYTHON_BIN:-python3}"
KERNEL_NAME="${KERNEL_NAME:-dronedetector-train}"
KERNEL_DISPLAY="${KERNEL_DISPLAY:-DroneDetector Train}"

log() {
  printf '[install_all_train] %s\n' "$*"
}

die() {
  printf '[install_all_train] ERROR: %s\n' "$*" >&2
  exit 1
}

run_as_root() {
  if [[ "${EUID}" -eq 0 ]]; then
    "$@"
  elif command -v sudo >/dev/null 2>&1; then
    sudo "$@"
  else
    die "Root privileges are required to install host packages."
  fi
}

preflight() {
  [[ -f "${PROJECT_ROOT}/deploy/requirements/nn_gpu_pinned.txt" ]] || die "Missing deploy/requirements/nn_gpu_pinned.txt"
  [[ -f "${PROJECT_ROOT}/requirements-train.txt" ]] || die "Missing train_scripts/requirements-train.txt"
  [[ -f "${PROJECT_ROOT}/torchsig/pyproject.toml" ]] || die "Missing local torchsig package"
  command -v "${PYTHON_BIN}" >/dev/null 2>&1 || die "${PYTHON_BIN} not found"
}

install_host_deps() {
  if ! command -v apt-get >/dev/null 2>&1; then
    log "apt-get not found; skipping host package installation"
    return
  fi

  log "Installing host packages"
  run_as_root apt-get update
  run_as_root env DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    python3 \
    python3-pip \
    python3-venv \
    python3-dev \
    build-essential \
    pkg-config \
    libgl1 \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev
}

setup_venv() {
  log "Creating/updating ${VENV_PATH}"
  if [[ ! -d "${VENV_PATH}" ]]; then
    "${PYTHON_BIN}" -m venv "${VENV_PATH}"
  fi

  # shellcheck disable=SC1090
  source "${VENV_PATH}/bin/activate"
  python -m pip install --upgrade pip setuptools wheel
}

install_python_deps() {
  log "Installing CUDA-enabled torch stack"
  python -m pip install -r "${PROJECT_ROOT}/deploy/requirements/nn_gpu_pinned.txt"

  log "Installing train notebook dependencies"
  python -m pip install -r "${PROJECT_ROOT}/train_scripts/requirements-train.txt"

  log "Installing local torchsig from repo"
  python -m pip install -e "${PROJECT_ROOT}/torchsig" --no-deps

  log "Registering Jupyter kernel ${KERNEL_NAME}"
  python -m ipykernel install --user --name "${KERNEL_NAME}" --display-name "${KERNEL_DISPLAY}"
}

verify_install() {
  log "Verifying imports"
  python - <<'PY'
import cv2
import ipykernel
import matplotlib
import mlconfig
import numpy
import pandas
import scipy
import sklearn
import torch
import torchvision
import torchsig

print('python_ok', True)
print('torch', torch.__version__)
print('torchvision', torchvision.__version__)
print('cuda_available', torch.cuda.is_available())
print('torch_cuda', torch.version.cuda)
print('torchsig', getattr(torchsig, '__version__', 'unknown'))
PY
}

main() {
  preflight

  if [[ "${1:-}" != "--skip-apt" ]]; then
    install_host_deps
  else
    log "Skipping host package installation (--skip-apt)"
  fi

  setup_venv
  install_python_deps
  verify_install
  log "SUCCESS: .venv-train is ready"
}

main "$@"