#!/bin/bash
# ═══════════════════════════════════════════════════════════════════════════════
# FIXED: start_nvidiam.sh (10.4.0.180)
# Safely boots Unsloth and vLLM-Omni on the same 16GB GPU without OOM crashes.
# ═══════════════════════════════════════════════════════════════════════════════
set -e

GREEN='\033[0;32m'; BLUE='\033[0;34m'; NC='\033[0m'
log() { echo -e "${BLUE}[NvidiaLLM]${NC} $1"; }
ok()  { echo -e "${GREEN}[NvidiaLLM] ✅${NC} $1"; }

log "Starting all NvidiaLLM services..."

if [ ! -d "$HOME/venvs/youtube-factory-nvidia" ]; then
    echo "❌ venv not found!"
    exit 1
fi

source ~/venvs/youtube-factory-nvidia/bin/activate

# 1. Brutally clean up any hung Python processes on these ports
log "Cleaning up old processes..."
pkill -f "unsloth.serve" || true
pkill -f "vllm_omni" || true
pkill -f "llm_server.py" || true
pkill -f "tts_server.py" || true
sleep 3

# 2. Start Unsloth LLM (GPU 0)
log "Starting Unsloth LLM on port 8001..."
CUDA_VISIBLE_DEVICES=0 \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
nohup python -m unsloth.serve \
    --model unsloth/Qwen3.5-7B-Instruct \
    --port 8001 \
    --max-model-len 8192 \
    > /tmp/unsloth-llm.log 2>&1 &

sleep 5
if curl -s http://localhost:8001/health > /dev/null 2>&1; then
    ok "Unsloth LLM running on port 8001"
else
    echo "⚠️  Unsloth still loading. Check: tail -f /tmp/unsloth-llm.log"
fi

# 3. Start vLLM-Omni (GPU 0) with strict VRAM limits
log "Starting vLLM-Omni TTS on port 8000..."
CUDA_VISIBLE_DEVICES=0 \
nohup python -m vllm_omni.entrypoints.cli.serve \
    --model step-audio-editx \
    --port 8000 \
    --dtype bfloat16 \
    --gpu-memory-utilization 0.55 \
    --enforce-eager \
    --max-model-len 4096 \
    --trust-remote-code \
    > /tmp/vllm-omni.log 2>&1 &

sleep 5
if curl -s http://localhost:8000/health > /dev/null 2>&1; then
    ok "vLLM-Omni TTS running on port 8000"
else
    echo "⚠️  vLLM-Omni still loading. Check: tail -f /tmp/vllm-omni.log"
fi

# 4. Verify ComfyUI
if curl -s http://localhost:8188/system_stats > /dev/null 2>&1; then
    ok "ComfyUI already running on port 8188"
else
    echo "⚠️  ComfyUI not detected on port 8188."
fi

echo ""
echo "📊 NvidiaLLM Status:"
echo "  Unsloth LLM:     http://10.4.0.180:8001"
echo "  vLLM-Omni TTS:   http://10.4.0.180:8000"
echo "  ComfyUI:         http://10.4.0.180:8188"