#!/usr/bin/env bash
# Install llama-server for OpenAI-compatible API on loopback :8080.
#
# Default: official Linux release asset (Ubuntu Vulkan — GPU without compile).
# CUDA source build only as last resort (or LLAMACPP_BUILD_CUDA=1).
# Pin: LLAMACPP_TAG=b10545  LLAMACPP_ASSET_URL=...  LLAMACPP_SHA256=...
set -euo pipefail

SWARM_USER="${SWARM_USER:-ubuntu}"
DATA_ROOT="/mnt/swarm_data"
LLAMA_ROOT="${DATA_ROOT}/llamacpp"
MODELS_DIR="${LLAMA_ROOT}/models"
BIN_DIR="${LLAMA_ROOT}/bin"
SRC_DIR="${LLAMA_ROOT}/src"
STAMP="${BIN_DIR}/.build-id"
UNIT="gpu-rent-llamacpp"
REPO="https://github.com/ggml-org/llama.cpp.git"
API_BASE="https://api.github.com/repos/ggml-org/llama.cpp"

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

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

mkdir -p "$MODELS_DIR" "$BIN_DIR"
chown -R "${SWARM_USER}:${SWARM_USER}" "$LLAMA_ROOT"

SERVER_BIN="${BIN_DIR}/llama-server"
LLAMACPP_TAG="${LLAMACPP_TAG:-}"
LLAMACPP_ASSET_URL="${LLAMACPP_ASSET_URL:-}"
LLAMACPP_SHA256="${LLAMACPP_SHA256:-}"
# 1 = force CUDA compile; 0 = never compile (Vulkan/CPU prebuilt only)
LLAMACPP_BUILD_CUDA="${LLAMACPP_BUILD_CUDA:-}"
# auto | cuda | vulkan — default auto: CUDA if nvcc already on VM, else Vulkan prebuilt
LLAMACPP_BACKEND="${LLAMACPP_BACKEND:-auto}"
LLAMACPP_FORCE_REINSTALL="${LLAMACPP_FORCE_REINSTALL:-}"
LLAMACPP_NGL="${LLAMACPP_NGL:-}"
LLAMACPP_CTX="${LLAMACPP_CTX:-}"
LLAMACPP_HOST="${LLAMACPP_HOST:-127.0.0.1}"
LLAMACPP_PORT="${LLAMACPP_PORT:-8080}"
LLAMACPP_EXTRA_ARGS="${LLAMACPP_EXTRA_ARGS:-}"

have_nvcc() {
  if command -v nvcc >/dev/null 2>&1; then
    return 0
  fi
  if [[ -x /usr/local/cuda/bin/nvcc ]]; then
    export PATH="/usr/local/cuda/bin:${PATH}"
    return 0
  fi
  return 1
}

# Prefer CUDA when toolkit already present (GPU images / prior up). Vulkan = fast no-compile.
want_cuda_build() {
  case "${LLAMACPP_BUILD_CUDA}" in
    1|yes|true) return 0 ;;
    0|no|false) return 1 ;;
  esac
  case "${LLAMACPP_BACKEND}" in
    cuda) return 0 ;;
    vulkan) return 1 ;;
    *)
      if have_nvcc; then
        return 0
      fi
      return 1
      ;;
  esac
}

if [[ "${LLAMACPP_FORCE_REINSTALL}" == "1" ]]; then
  log "LLAMACPP_FORCE_REINSTALL=1 — удаляю старый бинарь"
  rm -f "$SERVER_BIN" "$STAMP"
fi

# Upgrade path: previous default was Vulkan prebuilt; if nvcc is here, prefer CUDA.
if [[ -x "$SERVER_BIN" && -f "$STAMP" && "${LLAMACPP_BACKEND}" != "vulkan" && "${LLAMACPP_BUILD_CUDA}" != "0" ]]; then
  if grep -q '^asset:' "$STAMP" 2>/dev/null && want_cuda_build; then
    log "был Vulkan/CPU prebuilt, nvcc есть — пересобираю CUDA (лучше на NVIDIA)"
    rm -f "$SERVER_BIN" "$STAMP"
  fi
fi

# Prefer ubuntu CUDA (rare) → vulkan → cpu. Never Windows/macOS/cudart-only.
pick_linux_asset_url() {
  python3 -c '
import json,sys
data=json.load(sys.stdin)
assets=data.get("assets") or []
cands=[]
for a in assets:
    n=(a.get("name") or "").lower()
    u=a.get("browser_download_url") or ""
    if not (u.endswith(".zip") or u.endswith(".tar.gz")):
        continue
    if any(x in n for x in ("win","macos","android","darwin","xcframework","-ui.")):
        continue
    if "cudart" in n:
        continue
    score=0
    if "ubuntu" in n and "x64" in n and "cuda" in n:
        score=100
    elif "linux" in n and "cuda" in n:
        score=90
    elif "ubuntu" in n and "vulkan" in n and "x64" in n:
        score=50
    elif "ubuntu" in n and "x64" in n and not any(
        x in n for x in ("sycl","openvino","arm","s390","rocm")
    ):
        score=30
    elif "ubuntu" in n or "linux" in n:
        score=10
    if score:
        cands.append((score, u, n))
cands.sort(reverse=True)
print(cands[0][1] if cands else "")
'
}

resolve_release_tag() {
  # stdout = tag only (no log lines — callers capture via $())
  if [[ -n "$LLAMACPP_TAG" ]]; then
    echo "$LLAMACPP_TAG"
    return
  fi
  curl -fsSL "${API_BASE}/releases/latest" | python3 -c \
    'import json,sys; print(json.load(sys.stdin).get("tag_name") or "")'
}

cuda_architectures() {
  python3 - <<'PY'
import json
from pathlib import Path
p = Path("/mnt/swarm_data/.gpu-rent-gpu.json")
cap = "8.9"
if p.is_file():
    try:
        cap = str(json.loads(p.read_text()).get("compute_cap") or cap)
    except Exception:
        pass
parts = cap.split(".")
try:
    maj, mnr = int(parts[0]), int(parts[1]) if len(parts) > 1 else 0
    print(f"{maj}{mnr}")
except ValueError:
    print("89")
PY
}

ensure_build_deps() {
  export DEBIAN_FRONTEND=noninteractive
  apt-get install -y -qq \
    cmake build-essential git curl ca-certificates \
    libcurl4-openssl-dev >/dev/null
  if command -v nvcc >/dev/null 2>&1; then
    return 0
  fi
  if [[ -x /usr/local/cuda/bin/nvcc ]]; then
    export PATH="/usr/local/cuda/bin:${PATH}"
    return 0
  fi
  log "ставлю nvidia-cuda-toolkit (нужен nvcc)…"
  apt-get install -y -qq nvidia-cuda-toolkit >/dev/null
  if command -v nvcc >/dev/null 2>&1; then
    return 0
  fi
  if [[ -x /usr/local/cuda/bin/nvcc ]]; then
    export PATH="/usr/local/cuda/bin:${PATH}"
    return 0
  fi
  return 1
}

ensure_vulkan_runtime() {
  if ldconfig -p 2>/dev/null | grep -q 'libvulkan\.so'; then
    return 0
  fi
  export DEBIAN_FRONTEND=noninteractive
  log "ставлю libvulkan1 (для Ubuntu Vulkan prebuilt)…"
  apt-get install -y -qq libvulkan1 mesa-vulkan-drivers >/dev/null 2>&1 || \
    apt-get install -y -qq libvulkan1 >/dev/null 2>&1 || true
}

install_from_archive_url() {
  local url="$1"
  local tmp kind
  tmp="$(mktemp -d)"
  (
    cd "$tmp"
    log "скачиваю prebuilt: $url"
    curl -fL --progress-bar "$url" -o pkg.bin
    if [[ -n "$LLAMACPP_SHA256" ]]; then
      echo "${LLAMACPP_SHA256}  pkg.bin" | sha256sum -c -
    else
      log "WARN: LLAMACPP_SHA256 не задан — checksum skip"
    fi
    mkdir -p out
    # Do NOT grep -i zip — that matches "gzip" and breaks .tar.gz.
    kind="$(file -b pkg.bin 2>/dev/null || true)"
    case "$url" in
      *.zip)
        apt-get install -y -qq unzip >/dev/null 2>&1 || true
        unzip -qo pkg.bin -d out
        ;;
      *)
        if [[ "$kind" == Zip\ archive* ]] || [[ "$kind" == *"Zip archive"* ]]; then
          apt-get install -y -qq unzip >/dev/null 2>&1 || true
          unzip -qo pkg.bin -d out
        else
          tar -xaf pkg.bin -C out 2>/dev/null \
            || tar -xzf pkg.bin -C out 2>/dev/null \
            || tar -xf pkg.bin -C out
        fi
        ;;
    esac
    local found
    found="$(find out -type f -name 'llama-server' | head -n1 || true)"
    if [[ -z "$found" ]]; then
      found="$(find out -type f -name 'server' | head -n1 || true)"
    fi
    if [[ -z "$found" ]]; then
      log "в архиве нет llama-server (file says: ${kind:-unknown})"
      exit 1
    fi
    install -m 755 "$found" "$SERVER_BIN"
    # Shared libs next to binary (release tarballs ship .so alongside).
    find out -type f \( -name '*.so' -o -name '*.so.*' \) -print0 2>/dev/null \
      | while IFS= read -r -d '' so; do
          install -m 755 "$so" "${BIN_DIR}/$(basename "$so")"
        done
    chown -R "${SWARM_USER}:${SWARM_USER}" "$BIN_DIR"
  )
  local rc=$?
  rm -rf "$tmp"
  return "$rc"
}

install_linux_release() {
  local tag="$1"
  local api url
  api="${API_BASE}/releases/tags/${tag}"
  url="$(curl -fsSL "$api" | pick_linux_asset_url)"
  if [[ -z "$url" ]]; then
    log "в release ${tag} нет Linux-ассета"
    return 1
  fi
  if [[ "$url" == *vulkan* ]]; then
    log "беру Ubuntu Vulkan prebuilt (GPU без compile; CUDA-сборка — LLAMACPP_BUILD_CUDA=1)"
    ensure_vulkan_runtime
  elif [[ "$url" == *cuda* ]]; then
    log "беру Linux CUDA prebuilt"
  else
    log "WARN: Linux prebuilt без GPU backend (CPU) — ${url##*/}"
  fi
  if ! install_from_archive_url "$url"; then
    return 1
  fi
  echo "asset:${tag}" >"$STAMP"
  chown "${SWARM_USER}:${SWARM_USER}" "$STAMP"
}

# Quiet CUDA build: no cmake spam; heartbeat every 30s with last %.
build_cuda_from_source() {
  local tag="$1"
  local arch build_log pid pct line
  arch="$(cuda_architectures)"
  build_log="${LLAMA_ROOT}/build-cuda.log"
  log "крайний случай: сборка CUDA из исходников (tag=${tag}, arch=${arch}, 5–15 мин)…"
  log "полный лог: ${build_log}"
  if ! ensure_build_deps; then
    log "нет nvcc — CUDA-сборку пропускаем"
    return 1
  fi
  log "nvcc $(nvcc --version 2>/dev/null | tail -n1 || echo '?')"

  mkdir -p "$SRC_DIR"
  export GIT_TERMINAL_PROMPT=0
  if [[ -d "${SRC_DIR}/.git" ]]; then
    git -C "$SRC_DIR" -c advice.detachedHead=false fetch --depth 1 origin tag "$tag" 2>>"$build_log" || true
    if ! git -C "$SRC_DIR" -c advice.detachedHead=false checkout -f "$tag" >>"$build_log" 2>&1; then
      rm -rf "$SRC_DIR"
      git -c advice.detachedHead=false clone --depth 1 --branch "$tag" "$REPO" "$SRC_DIR" >>"$build_log" 2>&1
    fi
  else
    rm -rf "$SRC_DIR"
    git -c advice.detachedHead=false clone --depth 1 --branch "$tag" "$REPO" "$SRC_DIR" >>"$build_log" 2>&1
  fi

  cmake -S "$SRC_DIR" -B "${SRC_DIR}/build" \
    -DCMAKE_BUILD_TYPE=Release \
    -DGGML_CUDA=ON \
    -DCMAKE_CUDA_ARCHITECTURES="${arch}" \
    -DLLAMA_BUILD_SERVER=ON \
    -DLLAMA_BUILD_UI=OFF \
    -DLLAMA_USE_PREBUILT_UI=OFF \
    -DGGML_CCACHE=OFF \
    >>"$build_log" 2>&1

  # Background build + heartbeat (keeps SSH stream alive without 200 cmake lines).
  cmake --build "${SRC_DIR}/build" -j"$(nproc)" --target llama-server \
    >>"$build_log" 2>&1 &
  pid=$!
  while kill -0 "$pid" 2>/dev/null; do
    pct="$(grep -oE '\[[[:space:]]*[0-9]+%\]' "$build_log" 2>/dev/null | tail -n1 || true)"
    line="$(grep -E 'Building CUDA|Built target|Linking' "$build_log" 2>/dev/null | tail -n1 || true)"
    if [[ -n "$pct" ]]; then
      log "сборка CUDA ещё идёт… ${pct}${line:+ · ${line}}"
    else
      log "сборка CUDA ещё идёт… (cmake/nvcc, см. build.log)"
    fi
    sleep 30
  done
  if ! wait "$pid"; then
    log "сборка упала — хвост ${build_log}:"
    tail -n 40 "$build_log" >&2 || true
    return 1
  fi

  local built="${SRC_DIR}/build/bin/llama-server"
  if [[ ! -x "$built" ]]; then
    log "сборка не дала ${built}"
    return 1
  fi
  install -m 755 "$built" "$SERVER_BIN"
  # CUDA build may need libs from build/bin
  find "${SRC_DIR}/build/bin" -maxdepth 1 -type f \( -name '*.so' -o -name '*.so.*' \) -print0 2>/dev/null \
    | while IFS= read -r -d '' so; do
        install -m 755 "$so" "${BIN_DIR}/$(basename "$so")"
      done
  chown -R "${SWARM_USER}:${SWARM_USER}" "$BIN_DIR"
  echo "cuda:${tag}:${arch}" >"$STAMP"
  chown "${SWARM_USER}:${SWARM_USER}" "$STAMP"
  log "CUDA binary → ${SERVER_BIN}"
  return 0
}

normalize_tag() {
  printf '%s' "$1" | tr -d '\r' | head -n1 | awk 'NF{print; exit}'
}

if [[ -x "$SERVER_BIN" ]]; then
  log "llama-server уже есть: ${SERVER_BIN}"
else
  if [[ -n "$LLAMACPP_ASSET_URL" ]]; then
    log "скачиваю по LLAMACPP_ASSET_URL…"
    install_from_archive_url "$LLAMACPP_ASSET_URL"
    echo "asset-url" >"$STAMP"
    chown "${SWARM_USER}:${SWARM_USER}" "$STAMP"
  else
    if [[ -z "$LLAMACPP_TAG" ]]; then
      log "WARN: LLAMACPP_TAG не задан — latest (см. docs/llm.md)"
    fi
    tag="$(normalize_tag "$(resolve_release_tag)")"
    if [[ -z "$tag" || "$tag" == *" "* || "$tag" == *"["* ]]; then
      log "не удалось определить release tag (got: ${tag:-empty})"
      exit 1
    fi

    installed=0
    if want_cuda_build; then
      log "backend: CUDA (nvcc есть или LLAMACPP_BACKEND/BUILD_CUDA) — сборка, Vulkan только если упадёт"
      if build_cuda_from_source "$tag"; then
        installed=1
      else
        log "CUDA-сборка не вышла — fallback на Linux prebuilt (Vulkan/CPU)"
        if install_linux_release "$tag"; then
          installed=1
        fi
      fi
    else
      log "backend: Linux prebuilt (нет nvcc / LLAMACPP_BACKEND=vulkan) — без compile"
      if install_linux_release "$tag"; then
        installed=1
      else
        log "prebuilt не вышел — крайний случай: CUDA из исходников"
        if build_cuda_from_source "$tag"; then
          installed=1
        fi
      fi
    fi
    if [[ "$installed" != "1" ]]; then
      log "не удалось поставить llama-server"
      exit 1
    fi
  fi
fi

if [[ ! -x "$SERVER_BIN" ]]; then
  log "нет исполняемого ${SERVER_BIN}"
  exit 1
fi

# Prefer a weights GGUF (skip mmproj), then attach --mmproj if present.
MODEL_ARG=""
MMPROJ_ARG=""
FIRST_GGUF="$(
  find "$MODELS_DIR" -type f \( -name '*.gguf' -o -name '*.GGUF' \) \
    ! -iname '*mmproj*' 2>/dev/null | head -n1 || true
)"
MMPROJ_GGUF="$(
  find "$MODELS_DIR" -type f \( -iname '*mmproj*.gguf' -o -iname '*mmproj*.GGUF' \) \
    2>/dev/null | head -n1 || true
)"
if [[ -n "$FIRST_GGUF" ]]; then
  MODEL_ARG="-m ${FIRST_GGUF}"
  log "модель ${FIRST_GGUF}"
else
  log "нет GGUF в ${MODELS_DIR} — положи файл вручную и systemctl restart ${UNIT}"
fi
if [[ -n "$MMPROJ_GGUF" ]]; then
  MMPROJ_ARG="--mmproj ${MMPROJ_GGUF}"
  log "mmproj ${MMPROJ_GGUF}"
fi

# GPU layers: share card with Swarm — full offload on mid+, leave headroom on low.
NGL=99
CTX=8192
if [[ -f "${DATA_ROOT}/.gpu-rent-gpu.json" ]]; then
  eval "$(python3 - <<'PY'
import json
from pathlib import Path
gpu=json.loads(Path("/mnt/swarm_data/.gpu-rent-gpu.json").read_text())
vram=int(gpu.get("vram_mib") or 0)
gib=vram/1024.0
if gib < 16:
    print("NGL=40"); print("CTX=4096")
elif gib < 24:
    print("NGL=99"); print("CTX=8192")
elif gib < 48:
    print("NGL=99"); print("CTX=16384")
else:
    print("NGL=99"); print("CTX=32768")
PY
)" || true
fi
# Explicit overrides from gpu-rent.vars / .env (forwarded by provision).
if [[ -n "$LLAMACPP_NGL" ]]; then
  NGL="$LLAMACPP_NGL"
fi
if [[ -n "$LLAMACPP_CTX" ]]; then
  CTX="$LLAMACPP_CTX"
fi
log "llama.cpp -ngl ${NGL} -c ${CTX} host=${LLAMACPP_HOST} port=${LLAMACPP_PORT}${LLAMACPP_EXTRA_ARGS:+ extra=${LLAMACPP_EXTRA_ARGS}}"

cat >/etc/systemd/system/${UNIT}.service <<EOF
[Unit]
Description=gpu-rent llama.cpp server (loopback, GPU-tuned)
After=network-online.target local-fs.target
Wants=network-online.target

[Service]
Type=simple
User=${SWARM_USER}
Group=${SWARM_USER}
WorkingDirectory=${LLAMA_ROOT}
Environment=LD_LIBRARY_PATH=${BIN_DIR}
ExecStart=${SERVER_BIN} ${MODEL_ARG} ${MMPROJ_ARG} --host ${LLAMACPP_HOST} --port ${LLAMACPP_PORT} -ngl ${NGL} -c ${CTX} ${LLAMACPP_EXTRA_ARGS}
Restart=on-failure
RestartSec=8

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable "$UNIT"
systemctl restart "$UNIT" || log "unit стартовал с ошибкой (часто нет GGUF) — проверь journalctl -u ${UNIT}"
log "ok — http://127.0.0.1:8080 models=${MODELS_DIR}"