#!/usr/bin/env bash
# Install ComfyUI-HyMT2 dependencies into the ComfyUI venv.
#
# Hy-MT2 requires transformers >= 5.6.0. ⚠️ This conflicts with nodes that
# lock transformers==4.57.3 (Qwen3TTS, Sonic, certain VibeVoice forks).
# See INSTALL_NOTES.md for the rollback recipe.

set -euo pipefail

VENV=${COMFYUI_VENV:-/mnt/modelall/ComfyUI/.venv}
PYTHON="$VENV/bin/python"
TMPDIR=${TMPDIR:-/mnt/modelall/tmp}
mkdir -p "$TMPDIR"

if [ ! -x "$PYTHON" ]; then
  echo "❌ python not found at $PYTHON" >&2
  echo "   override the venv path: COMFYUI_VENV=/path/to/.venv $0" >&2
  exit 1
fi

echo "[install] venv=$VENV"
echo "[install] freezing current state to $TMPDIR/freeze_before_hymt2_$(date +%Y%m%d_%H%M).txt"
uv pip freeze --python "$PYTHON" \
  > "$TMPDIR/freeze_before_hymt2_$(date +%Y%m%d_%H%M).txt" || true

# Override file (v2.0 §2.3) — Py 3.13 needs newer wheels for these.
OVERRIDES="$TMPDIR/hymt2_overrides.txt"
cat > "$OVERRIDES" <<'EOF'
sentencepiece>=0.2.0
pyyaml>=6.0.2
huggingface_hub>=0.26
EOF

echo "[install] installing transformers>=5.6.0 + friends"
uv pip install --python "$PYTHON" \
  --override "$OVERRIDES" \
  --no-cache-dir \
  "transformers>=5.6.0" "tokenizers>=0.22.0" "accelerate>=0.30"

echo "[install] installing websocket-client (CLI only — not the node)"
uv pip install --python "$PYTHON" --no-cache-dir "websocket-client>=1.6"

echo "[install] smoke test — fail loudly if anything didn't install"
"$PYTHON" - <<'PY' || { echo "[install] ❌ SMOKE TEST FAILED — fix the missing packages above before restarting ComfyUI" >&2; exit 1; }
import sys
missing = []
versions = {}
for name in ("transformers", "tokenizers", "accelerate", "torch"):
    try:
        m = __import__(name)
        versions[name] = getattr(m, "__version__", "?")
    except ImportError as e:
        missing.append(f"{name} ({e})")
for n, v in versions.items():
    print(f"  {n}: {v}")
if missing:
    print(f"  MISSING: {', '.join(missing)}", file=sys.stderr)
    raise SystemExit(1)

# transformers >= 5.6 is required by Hy-MT2
import transformers
major = int(transformers.__version__.split(".")[0])
if major < 5:
    print(f"  ⚠️  transformers {transformers.__version__} < 5.6.0 — Hy-MT2 needs >= 5.6",
          file=sys.stderr)
    raise SystemExit(1)
print("  ✓ all required imports OK")
PY

echo "[install] DONE"
echo "[install] verify nodes registered after ComfyUI restart:"
echo "  curl -s http://localhost:8189/object_info | python3 -c \\"
echo "    \"import json,sys; d=json.load(sys.stdin); print('HyMT2Loader:', 'HyMT2Loader' in d, 'HyMT2Translate:', 'HyMT2Translate' in d)\""