#!/bin/bash
# Install LM Studio's headless daemon (llmster) for the agent user and start it at boot. Safe to re-run.
#
# The official installer downloads a 1.2 GB archive in one go and cannot resume, so on a flaky connection it
# often dies half way. This step performs the same actions (same download address, same sha512 check, same
# bootstrap command) but retries the download; if anything about that fails it falls back to the official script.
. "$(dirname "$0")/lib.sh"

LMS="$OWNER_HOME/.lmstudio/bin/lms"

install_manual() {
  tmp="$(mktemp -d /tmp/wxch-lms.XXXXXX)"; chmod 755 "$tmp"
  curl -fsSL --max-time 60 https://lmstudio.ai/install.sh -o "$tmp/install.sh" || return 1
  ver="$(sed -n 's/^APP_VERSION="\(.*\)"/\1/p' "$tmp/install.sh" | head -1)"
  down="$(sed -n 's/^ARTIFACT_DOWNLOAD_URL="\(.*\)"/\1/p' "$tmp/install.sh" | head -1)"
  [ -n "$ver" ] && [ -n "$down" ] || return 1
  rel="$ver-linux-$ARCH.full"
  ok=0
  for try in 1 2 3 4 5 6; do
    rm -f "$tmp/in.tar.gz"
    say "downloading llmster $ver (attempt $try of 6)"
    if curl -fL --connect-timeout 20 --speed-time 90 --speed-limit 20000 -o "$tmp/in.tar.gz" "https://$down/$rel.tar.gz"; then ok=1; break; fi
    sleep 5
  done
  [ "$ok" = 1 ] || return 1
  want="$(curl -fsSL "https://$down/$rel.sha512" | awk '{print $1}' | head -1)"
  have_sum="$(sha512sum "$tmp/in.tar.gz" | awk '{print $1}')"
  [ -n "$want" ] && [ "$want" = "$have_sum" ] || { warn "checksum mismatch for the llmster download"; return 1; }
  tar xf "$tmp/in.tar.gz" -C "$tmp" || return 1
  chown -R "$OWNER" "$tmp"
  as_owner env LMS_BOOTSTRAP_INSTALL_SH=1 "$tmp/llmster" bootstrap || return 1
  rm -rf "$tmp"
  return 0
}

if [ ! -x "$LMS" ]; then
  if ! install_manual; then
    warn "the resumable install did not work; trying the official installer"
    as_owner sh -c 'curl -fsSL https://lmstudio.ai/install.sh | LMS_NO_MODIFY_PATH=1 bash -s -- --no-modify-path' || die "the LM Studio install failed"
  fi
fi
[ -x "$LMS" ] || die "llmster is not installed ($LMS is missing)"

# keep the (large) model files in the model folder
mkdir -p "$MODELS_DIR/lmstudio"
chown "$OWNER" "$MODELS_DIR/lmstudio"
if [ ! -L "$OWNER_HOME/.lmstudio/models" ]; then
  if [ -d "$OWNER_HOME/.lmstudio/models" ] && [ -n "$(ls -A "$OWNER_HOME/.lmstudio/models" 2>/dev/null)" ]; then
    cp -a "$OWNER_HOME/.lmstudio/models/." "$MODELS_DIR/lmstudio/"
  fi
  rm -rf "$OWNER_HOME/.lmstudio/models"
  as_owner ln -s "$MODELS_DIR/lmstudio" "$OWNER_HOME/.lmstudio/models"
fi

# LM Studio rewrites its own lms program while the daemon starts; running it at that moment fails with
# "Text file busy". The start helper simply tries again a few times.
cat > /usr/local/bin/wxch-lms-start <<'HELPER'
#!/bin/sh
n=0
until "$1" server start --bind 127.0.0.1; do
  n=$((n+1)); [ "$n" -lt 10 ] || exit 1
  sleep 3
done
HELPER
chmod 755 /usr/local/bin/wxch-lms-start

# start at boot: the daemon plus its OpenAI-style server on loopback (models are loaded on demand by the agent)
cat > /etc/systemd/system/lmstudio.service <<EOF
[Unit]
Description=LM Studio (llmster) server
After=network-online.target
Wants=network-online.target
$(mountpoint -q "$MODELS_DIR" && echo "RequiresMountsFor=$MODELS_DIR")

[Service]
Type=oneshot
RemainAfterExit=yes
User=$OWNER
Environment="HOME=$OWNER_HOME"
ExecStartPre=$LMS daemon up
ExecStart=/usr/local/bin/wxch-lms-start $LMS
ExecStop=$LMS daemon down

[Install]
WantedBy=multi-user.target
EOF
systemctl daemon-reload
systemctl enable lmstudio.service >/dev/null 2>&1
systemctl restart lmstudio.service || { sleep 5; systemctl restart lmstudio.service; }

i=0
until curl -fsS http://127.0.0.1:1234/api/v1/models >/dev/null 2>&1; do
  i=$((i+1)); [ "$i" -lt 45 ] || die "the LM Studio server did not answer on 127.0.0.1:1234; see: journalctl -u lmstudio -n 50"
  sleep 1
done
say "LM Studio is up on 127.0.0.1:1234, models in $MODELS_DIR/lmstudio"