#!/bin/bash
#
# install_server.sh — idempotent server-side install for ComfyUI-VideoToText.
#
# Usage:
#     COMFYUI_DIR=/mnt/modelall/ComfyUI ./scripts/install_server.sh
#
# What it does:
#     1. Verifies ffmpeg is on PATH
#     2. Installs torchaudio into ComfyUI's venv (uv-managed)
#     3. Smoke-tests the imports
#
# What it does NOT do:
#     - Install MiMoASR (use that pack's own install_server.sh)
#     - Install rgthree-comfy (clone it manually if needed)
#     - Download model weights
#     - Restart ComfyUI (you do this after; see README Step 7)
#
set -euo pipefail

# ----- helpers -----
log()  { printf '\033[1;34m[+]\033[0m %s\n' "$*"; }
warn() { printf '\033[1;33m[!]\033[0m %s\n' "$*"; }
err()  { printf '\033[1;31m[x]\033[0m %s\n' "$*" >&2; }

# ----- 1. ffmpeg -----
log "Checking ffmpeg ..."
if ! command -v ffmpeg >/dev/null 2>&1; then
    err "ffmpeg not on PATH. Install with:  sudo apt install ffmpeg"
    exit 1
fi
ffmpeg -version | head -1

# ----- 2. ComfyUI venv -----
: "${COMFYUI_DIR:?Please set COMFYUI_DIR (e.g. /mnt/modelall/ComfyUI)}"
if [[ ! -d "$COMFYUI_DIR/.venv" ]]; then
    err "ComfyUI venv not found: $COMFYUI_DIR/.venv"
    exit 1
fi

log "Activating ComfyUI venv ..."
# shellcheck disable=SC1091
source "$COMFYUI_DIR/.venv/bin/activate"

# Pick uv if available; fall back to pip
if command -v uv >/dev/null 2>&1; then
    PIP="uv pip"
else
    PIP="pip"
fi

# ----- 3. torchaudio -----
log "Installing torchaudio ..."
$PIP install torchaudio

log "Smoke-test imports ..."
python - <<'PY'
import sys
ok = True
try:
    import torchaudio
    print(f"  torchaudio {torchaudio.__version__} OK")
except Exception as e:
    print(f"  torchaudio FAILED: {e}"); ok = False

try:
    import folder_paths  # ComfyUI core module
    print(f"  folder_paths OK ({folder_paths.__file__})")
except Exception as e:
    print(f"  folder_paths (ComfyUI core): cannot import outside ComfyUI runtime — that's normal, skipping")

sys.exit(0 if ok else 1)
PY

# ----- 4. final advice -----
cat <<'NEXT'

[+] torchaudio installed.

Next steps (manual):
  1. If not yet installed: clone ComfyUI-MiMoASR and rgthree-comfy
     $ cd "$COMFYUI_DIR/custom_nodes"
     $ git clone https://github.com/aadebuger/ComfyUI-MiMoASR.git
     $ git clone https://github.com/rgthree/rgthree-comfy.git

  2. Download MiMo weights (~24 GB):
     $ export HF_HOME=/mnt/modelall/hf_cache
     $ hf download XiaomiMiMo/MiMo-V2.5-ASR --local-dir /mnt/modelall/models/MiMo-V2.5-ASR
     $ hf download XiaomiMiMo/MiMo-Audio-Tokenizer --local-dir /mnt/modelall/models/MiMo-Audio-Tokenizer

  3. Clone MiMo upstream repo:
     $ git clone https://github.com/XiaomiMiMo/MiMo-V2.5-ASR.git /opt/repos/MiMo-V2.5-ASR

  4. Restart ComfyUI with required env vars:
     $ pkill -9 -f main.py; sleep 3
     $ export HF_HOME=/mnt/modelall/hf_cache
     $ export MIMO_ASR_REPO=/opt/repos/MiMo-V2.5-ASR
     $ cd "$COMFYUI_DIR" && source .venv/bin/activate
     $ nohup python main.py --listen 0.0.0.0 --port 8188 > ~/comfy_start.log 2>&1 &

  5. Verify 9 nodes registered (see README Step 8).

NEXT