#!/usr/bin/env bash
# Boots the Space: ComfyUI in the background, the studio in front.
#
# The UI comes up first on purpose. A Space counts as started once its port
# answers, and ComfyUI needs minutes to resolve weights and load a 4 GB model.
# The UI already polls for its backend and turns the header green by itself.
set -uo pipefail

COMFY=$HOME/ComfyUI
ENGINE_LOG=${ENGINE_LOG:-/tmp/comfyui.log}
CHECKPOINT=${YUE2_CHECKPOINT:-yue2_3b_int8_convrot.safetensors}
# The bf16 weights are what the Real-Audio adapter was built against, and what
# anyone with headroom should be picking anyway. Both are offered; the UI picks.
CHECKPOINT_ALT=${YUE2_CHECKPOINT_ALT:-yue2_3b_bf16.safetensors}
ENCODER=${SHEETSAGE_ENCODER:-sheetsage2_bf16.safetensors}
REAL_AUDIO_LORA=${YUE2_REAL_AUDIO_LORA:-nar_lora_joint_v9_comfyui.safetensors}
REAL_AUDIO_REPO=${YUE2_REAL_AUDIO_REPO:-Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4}
INST_LORA=${YUE2_INSTRUMENTAL_LORA:-ar_lora_inst_v3abc_comfyui.safetensors}
INST_REPO=${YUE2_INSTRUMENTAL_REPO:-Mothersuperior/YuE2-instrumental-cot-full-loras}

export ENGINE_LOG

# The hf CLI writes hints and warnings to stdout, so capturing its output in a
# command substitution yields banner text rather than a path. Ask the library
# directly instead, which prints the cached path and nothing else.
resolve() {
  python - "$1" "$2" <<'PY'
import sys
from huggingface_hub import hf_hub_download
print(hf_hub_download(repo_id=sys.argv[1], filename=sys.argv[2]))
PY
}

link_weight() {
  local repo=$1 path=$2 dest=$3
  local name src
  name=$(basename "$path")
  mkdir -p "$dest"
  if [ -e "$dest/$name" ]; then
    echo "[yue2] $name already present"
    return 0
  fi
  echo "[yue2] resolving $name"
  if ! src=$(resolve "$repo" "$path" 2>>"$ENGINE_LOG"); then
    echo "[yue2] ERROR could not resolve $name, see $ENGINE_LOG"
    return 1
  fi
  if [ ! -s "$src" ]; then
    echo "[yue2] ERROR resolved $name to an empty path: $src"
    return 1
  fi
  ln -sfn "$src" "$dest/$name"
  echo "[yue2] linked $name"
}

# Which image is actually running. Without this the only way to tell a stale
# container from a fresh one is to read a traceback's interpreter paths.
stamp() {
  {
    echo "[yue2] build ${BUILD_STAMP:-unstamped}"
    echo "[yue2] python $(python -c 'import sys; print(sys.version.split()[0])' 2>&1)"
    echo "[yue2] torch $(python -c 'import torch; print(torch.__version__)' 2>&1 | tail -1)"
  } | tee -a "$ENGINE_LOG"
}

start_comfy() {
  : > "$ENGINE_LOG"
  stamp

  link_weight Comfy-Org/YuE2 "checkpoints/$CHECKPOINT" "$COMFY/models/checkpoints" \
    || echo "[yue2] the checkpoint is missing, generation will fail" | tee -a "$ENGINE_LOG"
  link_weight Comfy-Org/YuE2 "checkpoints/$CHECKPOINT_ALT" "$COMFY/models/checkpoints" \
    || echo "[yue2] the bf16 checkpoint is missing, only the int8 one is selectable" | tee -a "$ENGINE_LOG"
  link_weight Comfy-Org/YuE2 "audio_encoders/$ENCODER" "$COMFY/models/audio_encoders" \
    || echo "[yue2] the transcriber is missing, covers will fail" | tee -a "$ENGINE_LOG"
  # Optional. Its absence greys the toggle out rather than breaking a render.
  link_weight "$REAL_AUDIO_REPO" "$REAL_AUDIO_LORA" "$COMFY/models/loras" \
    || echo "[yue2] the real-audio adapter is missing, that toggle stays off" | tee -a "$ENGINE_LOG"
  link_weight "$INST_REPO" "$INST_LORA" "$COMFY/models/loras" \
    || echo "[yue2] the instrumental adapter is missing, that toggle stays off" | tee -a "$ENGINE_LOG"

  cd "$COMFY"
  echo "[yue2] starting ComfyUI on 8188"
  # Tee so the Space log shows it live and the app can read the tail back.
  python main.py --listen 127.0.0.1 --port 8188 ${COMFY_EXTRA_ARGS:-} 2>&1 | tee -a "$ENGINE_LOG"
  echo "[yue2] ComfyUI exited. The studio stays up so the failure is visible." | tee -a "$ENGINE_LOG"
}

start_comfy &

cd "$HOME/app/.next/standalone"
echo "[yue2] starting the studio on ${PORT:-7860}"
export PORT="${PORT:-7860}" HOSTNAME=0.0.0.0
exec node server.js