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

PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
COMPOSE_FILE="${PROJECT_ROOT}/deploy/docker/docker-compose.yml"
SYSTEMD_TARGET_DIR="/etc/systemd/system"
RUN_USER="${SUDO_USER:-${USER}}"
RUN_GROUP="$(id -gn "${RUN_USER}")"

SOURCE_ARCHIVES=(
  "torchsig.tar.gz:torchsig:pyproject.toml"
)

SDR_UNITS=(
  dronedetector-sdr-433.service
  dronedetector-sdr-750.service
  dronedetector-sdr-1500.service
  dronedetector-sdr-3300.service
  dronedetector-sdr-4500.service
  dronedetector-sdr-5200.service
  dronedetector-sdr-5800.service
  dronedetector-sdr-915.service
  dronedetector-sdr-1200.service
  dronedetector-sdr-2400.service
)

CONFIGURED_SDR_UNITS=()

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

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

sdr_unit_env_key() {
  local unit="$1"
  local band="${unit#dronedetector-sdr-}"
  band="${band%.service}"
  printf 'hack_%s\n' "$band"
}

get_env_value() {
  local key="$1"
  awk -F= -v key="$key" '
    $1 == key {
      value = substr($0, index($0, "=") + 1)
      sub(/^[[:space:]]+/, "", value)
      sub(/[[:space:]]+$/, "", value)
      gsub(/^"/, "", value)
      gsub(/"$/, "", value)
      gsub(/^'\''/, "", value)
      gsub(/'\''$/, "", value)
      print value
      exit
    }
  ' "${PROJECT_ROOT}/.env"
}

populate_configured_sdr_units() {
  CONFIGURED_SDR_UNITS=()

  local unit env_key env_value
  for unit in "${SDR_UNITS[@]}"; do
    env_key="$(sdr_unit_env_key "$unit")"
    env_value="$(get_env_value "$env_key")"
    if [[ -n "$env_value" ]]; then
      CONFIGURED_SDR_UNITS+=("$unit")
    else
      log "Skipping ${unit}: ${env_key} is empty in .env"
    fi
  done

  if [[ "${#CONFIGURED_SDR_UNITS[@]}" -eq 0 ]]; then
    log "No SDR units are configured in .env"
  else
    log "Configured SDR units: ${CONFIGURED_SDR_UNITS[*]}"
  fi
}

print_failure_logs() {
  log "Collecting diagnostics..."
  systemctl --no-pager --full status dronedetector-compose.service || true
  for unit in "${SDR_UNITS[@]}"; do
    systemctl --no-pager --full status "$unit" || true
  done

  if command -v docker >/dev/null 2>&1; then
    docker compose -f "$COMPOSE_FILE" ps || true
    docker compose -f "$COMPOSE_FILE" logs --tail=150 dronedetector-server-to-master || true
    docker compose -f "$COMPOSE_FILE" logs --tail=150 dronedetector-nn-server || true
  fi

  journalctl -u dronedetector-compose.service -n 150 --no-pager || true
  for unit in "${SDR_UNITS[@]}"; do
    journalctl -u "$unit" -n 120 --no-pager || true
  done
}

trap 'rc=$?; if [[ $rc -ne 0 ]]; then print_failure_logs; fi' EXIT

require_root() {
  if [[ "${EUID}" -ne 0 ]]; then
    log "Switching to root via sudo..."
    exec sudo -E bash "$0" "$@"
  fi
}

extract_local_source_archives() {
  local spec archive_rel target_rel sentinel_rel
  local archive_path target_path sentinel_path

  for spec in "${SOURCE_ARCHIVES[@]}"; do
    IFS=':' read -r archive_rel target_rel sentinel_rel <<< "$spec"
    archive_path="${PROJECT_ROOT}/${archive_rel}"
    target_path="${PROJECT_ROOT}/${target_rel}"
    sentinel_path="${target_path}/${sentinel_rel}"

    if [[ -f "${sentinel_path}" ]]; then
      log "Vendored source already available: ${target_rel}"
      continue
    fi

    if [[ -e "${target_path}" ]]; then
      die "Found ${target_path}, but ${sentinel_rel} is missing. Remove or repair this directory, then rerun the installer."
    fi

    [[ -f "${archive_path}" ]] || die "Missing vendored source ${target_path} and archive ${archive_path}"
    command -v tar >/dev/null 2>&1 || die "tar is required to unpack ${archive_rel}"

    log "Extracting ${archive_rel} -> ${target_rel}"
    tar -xzf "${archive_path}" -C "${PROJECT_ROOT}"
    [[ -f "${sentinel_path}" ]] || die "Archive ${archive_path} did not unpack expected file ${sentinel_path}"
    chown -R "${RUN_USER}:${RUN_GROUP}" "${target_path}"
  done
}

preflight() {
  log "Preflight checks"
  [[ -f "${PROJECT_ROOT}/.env" ]] || die "Missing ${PROJECT_ROOT}/.env"
  [[ -f "${COMPOSE_FILE}" ]] || die "Missing ${COMPOSE_FILE}"

  if ! command -v apt-get >/dev/null 2>&1; then
    die "This installer currently supports Debian/Ubuntu only (apt-get required)."
  fi

  local free_mb
  free_mb="$(df -Pm "${PROJECT_ROOT}" | awk 'NR==2 {print $4}')"
  if [[ -z "$free_mb" || "$free_mb" -lt 10240 ]]; then
    die "At least 10 GB free disk space is required."
  fi

  if ! command -v nvidia-smi >/dev/null 2>&1; then
    die "nvidia-smi is required. GPU/NVIDIA driver is not available on host."
  fi

  log "Preflight OK"
}

install_host_non_python_deps() {
  log "Installing host non-python dependencies"
  apt-get update
  DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    ca-certificates \
    curl \
    gnupg \
    lsb-release \
    jq \
    git \
    python3 \
    python3-pip \
    python3-venv \
    build-essential \
    pkg-config \
    libusb-1.0-0 \
    libusb-1.0-0-dev \
    hackrf \
    gnuradio \
    gr-osmosdr
}

setup_sdr_python_env() {
  log "Setting up SDR python environment"
  local venv_path="${PROJECT_ROOT}/.venv-sdr"

  if [[ ! -d "$venv_path" ]]; then
    python3 -m venv --system-site-packages "$venv_path"
  fi

  "$venv_path/bin/pip" install --upgrade pip
  "$venv_path/bin/pip" install -r "${PROJECT_ROOT}/deploy/requirements/sdr_host.txt"

  chown -R "${RUN_USER}:${RUN_GROUP}" "$venv_path"
}

install_docker_if_needed() {
  if command -v docker >/dev/null 2>&1; then
    log "Docker already installed"
    return
  fi

  log "Installing Docker Engine"
  . /etc/os-release
  local distro_id="${ID}"
  if [[ "$distro_id" != "ubuntu" && "$distro_id" != "debian" ]]; then
    die "Unsupported distro for Docker auto-install: ${distro_id}"
  fi

  install -m 0755 -d /etc/apt/keyrings
  curl -fsSL "https://download.docker.com/linux/${distro_id}/gpg" | gpg --dearmor -o /etc/apt/keyrings/docker.gpg
  chmod a+r /etc/apt/keyrings/docker.gpg

  echo \
    "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/${distro_id} \
    ${VERSION_CODENAME} stable" | tee /etc/apt/sources.list.d/docker.list >/dev/null
  apt-get update
  DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin

  systemctl enable --now docker
}

install_nvidia_container_toolkit() {
  log "Installing/Configuring NVIDIA Container Toolkit"

  if ! dpkg -s nvidia-container-toolkit >/dev/null 2>&1; then
    curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | \
      gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

    curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
      sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
      tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

    apt-get update
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends nvidia-container-toolkit
  fi

  nvidia-ctk runtime configure --runtime=docker
  systemctl restart docker
}

build_and_run_compose() {
  log "Building and starting Docker services"
  docker compose -f "$COMPOSE_FILE" up -d --build
}

install_systemd_units() {
  log "Installing systemd units"

  local src dst
  for src in "${PROJECT_ROOT}"/deploy/systemd/*.service; do
    dst="${SYSTEMD_TARGET_DIR}/$(basename "$src")"
    sed \
      -e "s|__PROJECT_ROOT__|${PROJECT_ROOT}|g" \
      -e "s|__RUN_USER__|${RUN_USER}|g" \
      -e "s|__RUN_GROUP__|${RUN_GROUP}|g" \
      "$src" > "$dst"
  done

  systemctl daemon-reload
  systemctl enable dronedetector-compose.service
  systemctl restart dronedetector-compose.service

  for unit in "${CONFIGURED_SDR_UNITS[@]}"; do
    systemctl enable "$unit"
    systemctl restart "$unit"
  done
}

wait_for_systemd_active() {
  local unit="$1"
  local timeout_seconds="${2:-60}"
  local i

  for ((i=0; i<timeout_seconds; i++)); do
    if systemctl is-active --quiet "$unit"; then
      return 0
    fi
    sleep 1
  done

  return 1
}

verify_installation() {
  log "Verifying services"

  wait_for_systemd_active dronedetector-compose.service 30 || die "dronedetector-compose.service is not active"
  for unit in "${CONFIGURED_SDR_UNITS[@]}"; do
    wait_for_systemd_active "$unit" 45 || die "$unit is not active"
  done

  docker compose -f "$COMPOSE_FILE" ps

  local running_services
  running_services="$(docker compose -f "$COMPOSE_FILE" ps --status running --services || true)"

  printf '%s\n' "$running_services" | grep -Fxq "dronedetector-server-to-master" || die "server_to_master is not running"
  printf '%s\n' "$running_services" | grep -Fxq "dronedetector-telemetry-server" || die "telemetry_server is not running"

  log "Verification completed"
}

main() {
  require_root "$@"

  log "Project root: ${PROJECT_ROOT}"
  log "Runtime user: ${RUN_USER}:${RUN_GROUP}"

  extract_local_source_archives
  preflight
  populate_configured_sdr_units
  install_host_non_python_deps
  setup_sdr_python_env
  install_docker_if_needed
  install_nvidia_container_toolkit
  build_and_run_compose
  install_systemd_units
  verify_installation

  log "SUCCESS: DroneDetector installation completed"
}

main "$@"