#!/bin/bash
#
# MuseTalk 1.5 - Master Installation Script
# Instala todos os componentes necessários para o servidor
#
# Vendors instalados de /workspace/Bar Workspace/vendors/
#

set -e

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

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

echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          MUSETALK 1.5 - COMPLETE INSTALLATION                ║"
echo "║          Instala todos os componentes do servidor            ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""
echo "Vendors directory: $VENDORS_DIR"
echo ""

# ============================================================
# 0. CHECK PREREQUISITES
# ============================================================
header "[0/4] Checking prerequisites..."

# Check for Python
if ! command -v python3 &> /dev/null; then
    error "Python3 not found"
    exit 1
fi
success "Python3 found: $(python3 --version)"

# Check for pip
if ! command -v pip3 &> /dev/null && ! command -v pip &> /dev/null; then
    error "pip not found"
    exit 1
fi
success "pip found"

# Check for CUDA
if python3 -c "import torch; print(torch.cuda.is_available())" 2>/dev/null | grep -q "True"; then
    GPU_NAME=$(python3 -c "import torch; print(torch.cuda.get_device_name(0))" 2>/dev/null)
    success "CUDA available: $GPU_NAME"
else
    warn "CUDA not available, using CPU (slower)"
fi

# Check for ffmpeg
if command -v ffmpeg &> /dev/null; then
    success "ffmpeg found"
else
    warn "ffmpeg not found, installing..."
    sudo apt-get update -qq && sudo apt-get install -y -qq ffmpeg
    success "ffmpeg installed"
fi

# ============================================================
# 1. INSTALL ESPEAK-NG
# ============================================================
header "[1/4] Installing espeak-ng (TTS ~15ms)..."

if [ -f "$VENDORS_DIR/espeak/install.sh" ]; then
    bash "$VENDORS_DIR/espeak/install.sh"
else
    error "espeak vendor not found at $VENDORS_DIR/espeak/"
    exit 1
fi

# ============================================================
# 2. INSTALL WHISPER
# ============================================================
header "[2/4] Installing Distil-Whisper PT-BR (STT ~47ms)..."

if [ -f "$VENDORS_DIR/whisper/install.sh" ]; then
    bash "$VENDORS_DIR/whisper/install.sh"
else
    error "Whisper vendor not found at $VENDORS_DIR/whisper/"
    exit 1
fi

# ============================================================
# 3. INSTALL VLLM
# ============================================================
header "[3/4] Installing vLLM + Gemma 3 1B (LLM ~100ms)..."

if [ -f "$VENDORS_DIR/vllm/install.sh" ]; then
    bash "$VENDORS_DIR/vllm/install.sh"
else
    error "vLLM vendor not found at $VENDORS_DIR/vllm/"
    exit 1
fi

# ============================================================
# 4. VERIFY MUSETALK MODELS
# ============================================================
header "[4/4] Verifying MuseTalk models..."

MODELS_DIR="$SCRIPT_DIR/models"

check_model() {
    local name=$1
    local path=$2
    if [ -d "$path" ] && [ "$(ls -A $path 2>/dev/null)" ]; then
        success "$name found"
    else
        warn "$name missing at $path"
    fi
}

check_model "musetalk" "$MODELS_DIR/musetalk"
check_model "musetalkV15" "$MODELS_DIR/musetalkV15"
check_model "sd-vae-ft-mse" "$MODELS_DIR/sd-vae-ft-mse"
check_model "whisper" "$MODELS_DIR/whisper"
check_model "dwpose" "$MODELS_DIR/dwpose"
check_model "face-parse-bisent" "$MODELS_DIR/face-parse-bisent"

# ============================================================
# INSTALLATION COMPLETE
# ============================================================
echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          INSTALLATION COMPLETE!                              ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""
echo "  Installed components:"
echo "    ✓ espeak-ng     (TTS ~15ms)"
echo "    ✓ Distil-Whisper PT-BR (STT ~47ms)"
echo "    ✓ vLLM + Gemma 3 1B    (LLM ~100ms)"
echo ""
echo "  To start all services:"
echo "    bash $SCRIPT_DIR/server/start.sh"
echo ""
echo "  Or start individually from Bar Workspace:"
echo "    bash \"$VENDORS_DIR/whisper/start.sh\" 8766"
echo "    bash \"$VENDORS_DIR/vllm/start.sh\" 8000"
echo ""