#!/bin/bash
#
# MuseTalk Server - Complete Startup Script
# Starts all services with warm-up to keep models in GPU memory
#
# Individual scripts can be run separately:
#   bash server/start_vllm.sh     # Start vLLM (Gemma 3 1B)
#   bash server/start_whisper.sh  # Start Distil-Whisper PT-BR (STT)
#   bash server/start_musetalk.sh # Start MuseTalk (avatar server)
#

set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
BASE_DIR="$(dirname "$SCRIPT_DIR")"
cd "$BASE_DIR"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

success() { echo -e "${GREEN}✓${NC} $1"; }
warn() { echo -e "${YELLOW}⚠${NC} $1"; }
error() { echo -e "${RED}✗${NC} $1"; }
info() { echo -e "  $1"; }
header() { echo -e "${CYAN}$1${NC}"; }

echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          MUSETALK SERVER - STARTING ALL SERVICES             ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""

# ============================================================
# 0. STOP EXISTING SERVICES
# ============================================================
header "[0/4] Stopping existing services..."

# Kill existing MuseTalk processes
if pgrep -f "main_fast" > /dev/null; then
    pkill -f "main_fast" 2>/dev/null || true
    sleep 2
fi

success "Cleanup done"

# ============================================================
# 1. START VLLM (Gemma 3 1B)
# ============================================================
echo ""
header "[1/4] Starting vLLM (Gemma 3 1B)..."

if bash "$SCRIPT_DIR/start_vllm.sh" 8000; then
    success "vLLM ready (port 8000)"
else
    error "vLLM failed to start"
    exit 1
fi

# ============================================================
# 2. START DISTIL-WHISPER PT-BR (STT)
# ============================================================
echo ""
header "[2/4] Starting Distil-Whisper PT-BR (STT)..."

if bash "$SCRIPT_DIR/start_whisper.sh" 8766; then
    success "Distil-Whisper ready (port 8766)"
else
    warn "Distil-Whisper failed to start (will use Groq API fallback)"
fi

# ============================================================
# 3. START MUSETALK SERVER
# ============================================================
echo ""
header "[3/4] Starting MuseTalk server..."

if bash "$SCRIPT_DIR/start_musetalk.sh" 3000; then
    success "MuseTalk ready (port 3000)"
else
    error "MuseTalk failed to start"
    exit 1
fi

# ============================================================
# 4. VERIFY ALL SERVICES
# ============================================================
echo ""
header "[4/4] Verifying all services..."

echo ""
echo "┌────────────────────────────────────────────────────────────────┐"
echo "│ Service          │ Port  │ Status   │ Latency                 │"
echo "├────────────────────────────────────────────────────────────────┤"

# Check vLLM
if curl -s http://localhost:8000/v1/models > /dev/null 2>&1; then
    echo "│ vLLM (Gemma 3)   │ 8000  │ ✓ Online │ ~100ms (LLM)            │"
else
    echo "│ vLLM (Gemma 3)   │ 8000  │ ✗ Offline│ -                       │"
fi

# Check Distil-Whisper
if curl -s "http://localhost:8766/health" 2>/dev/null | grep -q "ok"; then
    echo "│ Distil-Whisper   │ 8766  │ ✓ Online │ ~47ms (STT PT-BR)       │"
else
    echo "│ Distil-Whisper   │ 8766  │ ✗ Offline│ -                       │"
fi

# Check MuseTalk
if curl -s http://localhost:3000/api/health 2>/dev/null | grep -q "ok"; then
    echo "│ MuseTalk         │ 3000  │ ✓ Online │ Avatar ready            │"
else
    echo "│ MuseTalk         │ 3000  │ ✗ Offline│ -                       │"
fi

echo "└────────────────────────────────────────────────────────────────┘"

# ============================================================
# SUMMARY
# ============================================================
echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          ALL SERVICES STARTED!                               ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""
echo "  Web interface:  http://localhost:3000"
echo "  Demo mode:      http://localhost:3000/?demo=true"
echo "  WebRTC mode:    http://localhost:3000/rtc"
echo ""
echo "  Individual service scripts:"
echo "    bash server/start_vllm.sh     # vLLM (Gemma 3 1B)"
echo "    bash server/start_whisper.sh  # Distil-Whisper PT-BR (STT)"
echo "    bash server/start_musetalk.sh # MuseTalk (avatar server)"
echo ""
echo "  Logs:"
echo "    MuseTalk:       tail -f /tmp/server_fast.log"
echo "    Distil-Whisper: tail -f /tmp/whisper_distil.log"
echo "    vLLM:           tail -f /tmp/vllm.log"
echo ""
echo "  Benchmark:   python3 tests/benchmark_pipeline.py 5"
echo ""
echo "  Pipeline latency breakdown:"
echo "    STT (Distil-Whisper): ~47ms"
echo "    LLM (Gemma 3 1B):     ~100ms"
echo "    TTS (espeak):         ~15ms"
echo "    MuseTalk 1st frame:   ~480ms"
echo "    ─────────────────────────────"
echo "    Total:                ~642ms"
echo ""