#!/bin/bash
#
# dumont-talker - Start All Services
#

set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"
VENDORS_DIR="/workspace/Bar Workspace/vendors"

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"; }
header() { echo -e "\n${CYAN}$1${NC}"; }

check_service() {
    curl -s "http://localhost:$1$2" > /dev/null 2>&1
}

echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          DUMONT-TALKER - START SERVICES                      ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""

# ============================================================
# 1. WHISPER (STT)
# ============================================================
header "1. Whisper (STT) - Port 8766"

if check_service 8766 "/health"; then
    success "Already running"
else
    if [ -f "$VENDORS_DIR/whisper/start.sh" ]; then
        "$VENDORS_DIR/whisper/start.sh" 8766 &
        sleep 5
        check_service 8766 "/health" && success "Started" || warn "Starting..."
    else
        warn "Whisper start script not found"
    fi
fi

# ============================================================
# 2. VLLM (LLM)
# ============================================================
header "2. vLLM (LLM) - Port 8000"

if check_service 8000 "/v1/models"; then
    success "Already running"
else
    if [ -f "$VENDORS_DIR/vllm/start.sh" ]; then
        "$VENDORS_DIR/vllm/start.sh" 8000 &
        sleep 10
        check_service 8000 "/v1/models" && success "Started" || warn "Starting..."
    else
        warn "vLLM start script not found"
    fi
fi

# ============================================================
# 3. MUSETALK (Avatar)
# ============================================================
header "3. MuseTalk (Avatar) - Port 3000"

if check_service 3000 "/api/health"; then
    success "Already running"
else
    export FFMPEG_PATH=/usr/bin
    export HF_HOME=/tmp/hf_home
    export PYTHONPATH="$PROJECT_DIR/server:$PYTHONPATH"
    
    cd "$PROJECT_DIR/server"
    nohup python3 main.py 3000 > /tmp/dumont-talker.log 2>&1 &
    echo $! > /tmp/dumont-talker.pid
    
    echo "  Waiting for MuseTalk to load models..."
    for i in {1..30}; do
        if check_service 3000 "/api/health"; then
            success "Started"
            break
        fi
        sleep 4
    done
fi

# ============================================================
# STATUS
# ============================================================
header "SERVICE STATUS"
echo ""

printf "  %-12s %-8s %s\n" "Service" "Port" "Status"
echo "  ──────────────────────────────────"

for svc in "Whisper:8766:/health" "vLLM:8000:/v1/models" "MuseTalk:3000:/api/health"; do
    IFS=':' read -r name port endpoint <<< "$svc"
    if check_service "$port" "$endpoint"; then
        printf "  %-12s %-8s ${GREEN}✓ Online${NC}\n" "$name" "$port"
    else
        printf "  %-12s %-8s ${YELLOW}○ Offline${NC}\n" "$name" "$port"
    fi
done

# ============================================================
# WARMUP (5 inferences per service)
# ============================================================
header "WARMUP (5 inferences each)"

# Warmup Whisper
if check_service 8766 "/health"; then
    echo "  Warming up Whisper..."
    DEMO_AUDIO="$PROJECT_DIR/web/demo_audio.webm"
    if [ -f "$DEMO_AUDIO" ]; then
        for i in {1..5}; do
            curl -s -X POST "http://localhost:8766/inference" \
                -F "file=@$DEMO_AUDIO" \
                -F "response_format=json" \
                -F "language=pt" > /dev/null 2>&1
        done
        success "Whisper warmed up"
    fi
fi

# Warmup vLLM
if check_service 8000 "/v1/models"; then
    echo "  Warming up vLLM..."
    for i in {1..5}; do
        curl -s -X POST "http://localhost:8000/v1/chat/completions" \
            -H "Content-Type: application/json" \
            -d '{"model":"google/gemma-3-1b-it","messages":[{"role":"user","content":"Olá"}],"max_tokens":20}' > /dev/null 2>&1
    done
    success "vLLM warmed up"
fi

# Warmup MuseTalk
if check_service 3000 "/api/health"; then
    echo "  Warming up MuseTalk..."
    DEMO_AUDIO="$PROJECT_DIR/web/demo_audio.webm"
    if [ -f "$DEMO_AUDIO" ]; then
        for i in {1..3}; do
            curl -s -X POST "http://localhost:3000/api/conversation" \
                -F "audio=@$DEMO_AUDIO" \
                -F "resolution=256" \
                -F "batch_size=8" > /dev/null 2>&1
        done
        success "MuseTalk warmed up"
    fi
fi

echo ""
echo "  Endpoints:"
echo "    • Web UI:    http://localhost:3000"
echo "    • Logs:      /tmp/dumont-talker.log"
echo ""
echo "  ✅ All services ready!"
echo ""