#!/usr/bin/env bash
# Idempotent gpu-rent first-boot on the VM. No Docker. Do not apt-upgrade the kernel.
set -euo pipefail

SWARM_USER="${SWARM_USER:-ubuntu}"
SWARM_ROOT="/opt/swarmui"
DATA_ROOT="/mnt/swarm_data"
MARKER_DATA="${DATA_ROOT}/.gpu-rent-ready"
MARKER_BOOT="${SWARM_ROOT}/.gpu-rent-bootstrapped"
SWARM_REPO="https://github.com/mcmonkeyprojects/SwarmUI.git"

log() { echo "[gpu-rent] $*" >&2; }

if [[ "$(id -u)" -ne 0 ]]; then
  echo "нужен root (sudo -n bash $0)" >&2
  exit 1
fi

export DEBIAN_FRONTEND=noninteractive

pick_data_disk() {
  local root_src root_pk name type
  root_src="$(findmnt -n -o SOURCE / || true)"
  root_pk="$(lsblk -no PKNAME "$root_src" 2>/dev/null | head -n1 || true)"
  if [[ -z "$root_pk" && -n "$root_src" ]]; then
    root_pk="$(lsblk -no NAME "$root_src" 2>/dev/null | head -n1 | sed 's/[0-9]*$//' || true)"
  fi
  while read -r name type; do
    [[ "$type" == "disk" ]] || continue
    [[ -n "$root_pk" && "$name" == "$root_pk" ]] && continue
    if findmnt "/dev/${name}" >/dev/null 2>&1; then
      continue
    fi
    echo "/dev/${name}"
    return 0
  done < <(lsblk -dn -o NAME,TYPE)
  return 1
}

ensure_data_mount() {
  mkdir -p "$DATA_ROOT"
  if findmnt "$DATA_ROOT" >/dev/null 2>&1; then
    log "data volume уже на ${DATA_ROOT}"
    return 0
  fi
  local disk fstype uuid
  disk="$(pick_data_disk)" || {
    log "нет второго диска — проверь BDM data volume"
    exit 1
  }
  fstype="$(blkid -s TYPE -o value "$disk" 2>/dev/null || true)"
  if [[ -z "$fstype" ]]; then
    if [[ -e "$MARKER_DATA" ]]; then
      log "маркер есть, а FS на ${disk} нет — не mkfs, разбери вручную"
      exit 1
    fi
    log "mkfs.ext4 ${disk} (пустой data volume)"
    mkfs.ext4 -F -L swarm-data "$disk"
  fi
  mount "$disk" "$DATA_ROOT"
  uuid="$(blkid -s UUID -o value "$disk")"
  if [[ -n "$uuid" ]] && ! grep -q "UUID=${uuid}" /etc/fstab; then
    echo "UUID=${uuid} ${DATA_ROOT} ext4 defaults,nofail 0 2" >> /etc/fstab
  fi
  log "смонтирован ${disk} -> ${DATA_ROOT}"
}

ensure_bind() {
  local src="$1" dst="$2"
  mkdir -p "$src" "$dst"
  if ! findmnt "$dst" >/dev/null 2>&1; then
    mount --bind "$src" "$dst"
  fi
  if ! grep -Fq " ${dst} " /etc/fstab; then
    echo "${src} ${dst} none bind,nofail 0 0" >> /etc/fstab
  fi
}

log "пакеты (без upgrade ядра)"
# Light: skip apt when Swarm already bootstrapped, or llm-only data disk already ready.
_light_ok=0
if [[ "${GPU_RENT_BOOTSTRAP_LIGHT:-0}" == "1" ]]; then
  if [[ -f "$MARKER_BOOT" ]]; then
    _light_ok=1
  elif [[ "${GPU_RENT_SKIP_SWARMUI:-0}" == "1" && -f "$MARKER_DATA" ]]; then
    _light_ok=1
  fi
fi
if [[ "$_light_ok" == "1" ]]; then
  log "light bootstrap — пропускаем apt-get (маркер уже есть)"
else
  if [[ "${GPU_RENT_BOOTSTRAP_LIGHT:-0}" == "1" ]]; then
    log "light запрошен, но маркера нет — полный apt"
  fi
  apt-get update -qq
  apt-get install -y -qq git python3 python3-venv python3-pip python3-full curl ca-certificates
fi

ensure_data_mount

mkdir -p \
  "${DATA_ROOT}/Models" \
  "${DATA_ROOT}/Output" \
  "${DATA_ROOT}/Data" \
  "${DATA_ROOT}/Data/Autocompletions" \
  "${DATA_ROOT}/Data/Wildcards" \
  "${DATA_ROOT}/dlbackend" \
  "${DATA_ROOT}/Extensions" \
  "${DATA_ROOT}/DLNodes" \
  "${DATA_ROOT}/CustomWorkflows" \
  "${DATA_ROOT}/ollama"

# LLM-only: data disk + tools, no SwarmUI clone / unit.
if [[ "${GPU_RENT_SKIP_SWARMUI:-0}" == "1" ]]; then
  chown -R "${SWARM_USER}:${SWARM_USER}" "$DATA_ROOT"
  date -u +"%Y-%m-%dT%H:%M:%SZ" >"$MARKER_DATA"
  date -u +"%Y-%m-%dT%H:%M:%SZ" >"${DATA_ROOT}/.gpu-rent-llm-only"
  chown "${SWARM_USER}:${SWARM_USER}" "$MARKER_DATA" "${DATA_ROOT}/.gpu-rent-llm-only"
  systemctl stop swarmui 2>/dev/null || true
  systemctl disable swarmui 2>/dev/null || true
  if command -v nvidia-smi >/dev/null 2>&1; then
    nvidia-smi -L || true
  fi
  log "bootstrap ok (llm-only; без SwarmUI)"
  exit 0
fi

rm -f "${DATA_ROOT}/.gpu-rent-llm-only" 2>/dev/null || true

if [[ ! -d "${SWARM_ROOT}/.git" ]]; then
  log "clone SwarmUI -> ${SWARM_ROOT}"
  mkdir -p "$(dirname "$SWARM_ROOT")"
  git clone --depth 1 "$SWARM_REPO" "$SWARM_ROOT"
elif [[ "${GPU_RENT_UPDATE_GIT:-1}" == "1" ]]; then
  log "обновляю SwarmUI в ${SWARM_ROOT}"
  branch="$(git -C "$SWARM_ROOT" remote show origin 2>/dev/null | sed -n '/HEAD branch/s/.*: //p' || true)"
  branch="${branch:-master}"
  # shallow clone: deepen tip of default branch
  git -C "$SWARM_ROOT" fetch --depth 1 origin "$branch" || git -C "$SWARM_ROOT" fetch --depth 1 origin
  if git -C "$SWARM_ROOT" rev-parse --verify -q "origin/${branch}" >/dev/null; then
    git -C "$SWARM_ROOT" checkout -B "$branch" "origin/${branch}"
    git -C "$SWARM_ROOT" reset --hard "origin/${branch}"
  else
    git -C "$SWARM_ROOT" pull --ff-only || true
  fi
  log "SwarmUI @ $(git -C "$SWARM_ROOT" rev-parse --short HEAD)"
else
  log "SwarmUI уже в ${SWARM_ROOT} (update off)"
fi

if [[ ! -x /usr/share/dotnet/dotnet && ! -x "/home/${SWARM_USER}/.dotnet/dotnet" ]]; then
  if [[ -f "${SWARM_ROOT}/launchtools/linux-dotnet-install.sh" ]]; then
    log "ставим .NET SDK (скрипт SwarmUI)"
    # SwarmUI's linux-dotnet-install.sh does `cd launchtools` and
    # `./launchtools/dotnet-install.sh` — must run from repo root.
    if ! sudo -u "$SWARM_USER" bash -c "cd \"$SWARM_ROOT\" && bash launchtools/linux-dotnet-install.sh"; then
      log "WARN: linux-dotnet-install.sh завершился с ошибкой"
    fi
  fi
  if [[ ! -x /usr/share/dotnet/dotnet && ! -x "/home/${SWARM_USER}/.dotnet/dotnet" ]]; then
    log "ERROR: .NET не установлен ($HOME/.dotnet или /usr/share/dotnet) — SwarmUI не стартует"
    exit 1
  fi
  log ".NET ok: $(sudo -u "$SWARM_USER" bash -lc 'command -v dotnet || echo /home/'"$SWARM_USER"'/.dotnet/dotnet')"
fi

mkdir -p \
  "${SWARM_ROOT}/Models" \
  "${SWARM_ROOT}/Output" \
  "${SWARM_ROOT}/Data" \
  "${SWARM_ROOT}/dlbackend" \
  "${SWARM_ROOT}/src/Extensions" \
  "${SWARM_ROOT}/src/BuiltinExtensions/ComfyUIBackend/DLNodes" \
  "${SWARM_ROOT}/src/BuiltinExtensions/ComfyUIBackend/CustomWorkflows"

ensure_bind "${DATA_ROOT}/Models" "${SWARM_ROOT}/Models"
ensure_bind "${DATA_ROOT}/Output" "${SWARM_ROOT}/Output"
ensure_bind "${DATA_ROOT}/Data" "${SWARM_ROOT}/Data"
ensure_bind "${DATA_ROOT}/dlbackend" "${SWARM_ROOT}/dlbackend"
ensure_bind "${DATA_ROOT}/Extensions" "${SWARM_ROOT}/src/Extensions"
ensure_bind "${DATA_ROOT}/DLNodes" "${SWARM_ROOT}/src/BuiltinExtensions/ComfyUIBackend/DLNodes"
ensure_bind "${DATA_ROOT}/CustomWorkflows" "${SWARM_ROOT}/src/BuiltinExtensions/ComfyUIBackend/CustomWorkflows"

chown -R "${SWARM_USER}:${SWARM_USER}" "$DATA_ROOT" "$SWARM_ROOT"

cat >/etc/systemd/system/swarmui.service <<EOF
[Unit]
Description=SwarmUI (gpu-rent)
After=network-online.target local-fs.target
Wants=network-online.target

[Service]
Type=simple
User=${SWARM_USER}
Group=${SWARM_USER}
WorkingDirectory=${SWARM_ROOT}
Environment=HOME=/home/${SWARM_USER}
Environment=DOTNET_ROOT=/home/${SWARM_USER}/.dotnet
Environment=DOTNET_CLI_HOME=/home/${SWARM_USER}
Environment=PATH=/home/${SWARM_USER}/.dotnet:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
ExecStart=${SWARM_ROOT}/launch-linux.sh --launch_mode none --host 127.0.0.1 --port 7801
Restart=on-failure
RestartSec=8
TimeoutStartSec=0

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable swarmui
# Не стартуем UI здесь: сначала extensions / autocomplete / Civitai / push.

date -u +"%Y-%m-%dT%H:%M:%SZ" >"$MARKER_DATA"
date -u +"%Y-%m-%dT%H:%M:%SZ" >"$MARKER_BOOT"
chown "${SWARM_USER}:${SWARM_USER}" "$MARKER_DATA" "$MARKER_BOOT"

if command -v nvidia-smi >/dev/null 2>&1; then
  nvidia-smi -L || true
fi
log "bootstrap ok (unit готов; swarmui старт после seed)"