#!/bin/bash
#
# MuseTalk Server - Complete Installation Script
# Installs: vLLM (Gemma 3 1B), Whisper, models, and all dependencies
#

set -e

echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          MUSETALK SERVER - INSTALLATION                      ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""

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

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

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

# ============================================================
# 1. SYSTEM DEPENDENCIES
# ============================================================
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "[1/6] Checking system dependencies..."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# Check for required commands
MISSING_DEPS=""

if ! command -v ffmpeg &> /dev/null; then
    MISSING_DEPS="$MISSING_DEPS ffmpeg"
fi

if ! command -v espeak-ng &> /dev/null; then
    MISSING_DEPS="$MISSING_DEPS espeak-ng"
fi

if ! command -v cmake &> /dev/null; then
    MISSING_DEPS="$MISSING_DEPS cmake"
fi

if [ -n "$MISSING_DEPS" ]; then
    warn "Missing dependencies:$MISSING_DEPS"
    info "Installing with apt..."
    sudo apt-get update -qq
    sudo apt-get install -y -qq ffmpeg espeak-ng cmake build-essential
    success "System dependencies installed"
else
    success "System dependencies already installed"
fi

# ============================================================
# 2. VLLM INSTALLATION (Gemma 3 1B)
# ============================================================
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "[2/6] Installing vLLM (Gemma 3 1B - fast inference)..."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

if python3 -c "import vllm" 2>/dev/null; then
    VLLM_VERSION=$(python3 -c "import vllm; print(vllm.__version__)" 2>/dev/null)
    success "vLLM already installed: v$VLLM_VERSION"
else
    info "Installing vLLM (this may take a few minutes)..."
    pip install vllm --quiet
    success "vLLM installed"
fi

# Pre-download Gemma 3 1B model
info "Pre-downloading Gemma 3 1B model..."
export HF_HOME=/tmp/hf_vllm
# HF_TOKEN should be set in environment or .env file
export HF_TOKEN="${HF_TOKEN:-}"
export HUGGING_FACE_HUB_TOKEN="${HF_TOKEN:-}"
mkdir -p /tmp/hf_vllm

python3 -c "
from huggingface_hub import snapshot_download
snapshot_download('google/gemma-3-1b-it', local_files_only=False)
print('Model downloaded')
" 2>/dev/null || warn "Model will be downloaded on first start"
success "Gemma 3 1B model ready"

# ============================================================
# 3. DISTIL-WHISPER PT-BR INSTALLATION (STT)
# ============================================================
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "[3/6] Installing Distil-Whisper PT-BR (STT)..."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# Install Python dependencies for Distil-Whisper
info "Installing transformers and dependencies..."
pip install transformers accelerate librosa --quiet

# Pre-download Distil-Whisper PT-BR model
info "Pre-downloading Distil-Whisper PT-BR model..."
export HF_HOME=/tmp/hf_whisper
mkdir -p /tmp/hf_whisper
python3 -c "
from huggingface_hub import snapshot_download
snapshot_download('freds0/distil-whisper-large-v3-ptbr', local_files_only=False)
print('Model downloaded')
" 2>/dev/null || warn "Model will be downloaded on first start"

success "Distil-Whisper PT-BR ready (~47ms latency)"

cd "$BASE_DIR"

# ============================================================
# 4. PYTHON DEPENDENCIES
# ============================================================
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "[4/6] Installing Python dependencies..."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# Check if key packages are installed
MISSING_PY=""

python3 -c "import fastapi" 2>/dev/null || MISSING_PY="$MISSING_PY fastapi"
python3 -c "import uvicorn" 2>/dev/null || MISSING_PY="$MISSING_PY uvicorn"
python3 -c "import httpx" 2>/dev/null || MISSING_PY="$MISSING_PY httpx"
python3 -c "import aiortc" 2>/dev/null || MISSING_PY="$MISSING_PY aiortc"
python3 -c "import av" 2>/dev/null || MISSING_PY="$MISSING_PY av"

if [ -n "$MISSING_PY" ]; then
    info "Installing:$MISSING_PY"
    pip install fastapi uvicorn httpx aiortc av aiofiles python-multipart --quiet
    success "Python dependencies installed"
else
    success "Python dependencies already installed"
fi

# ============================================================
# 5. MUSETALK MODELS
# ============================================================
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "[5/6] Downloading MuseTalk models..."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

mkdir -p "$BASE_DIR/models"
cd "$BASE_DIR"

# Download MuseTalk models from HuggingFace
info "Downloading MuseTalk models (this may take a while)..."
export HF_HOME=/tmp/hf_musetalk

python3 -c "
from huggingface_hub import snapshot_download
import os

models_dir = '$BASE_DIR/models'

# MuseTalk main model
print('Downloading MuseTalk model...')
snapshot_download('TMElyralab/MuseTalk', local_dir=f'{models_dir}/musetalk', local_dir_use_symlinks=False)

# SD VAE
print('Downloading SD VAE...')
snapshot_download('stabilityai/sd-vae-ft-mse', local_dir=f'{models_dir}/sd-vae-ft-mse', local_dir_use_symlinks=False)

# DWPose
print('Downloading DWPose...')
snapshot_download('yzd-v/DWPose', local_dir=f'{models_dir}/dwpose', local_dir_use_symlinks=False)

# Face parsing
print('Downloading face-parse-bisent...')
snapshot_download('jonathandinu/face-parsing', local_dir=f'{models_dir}/face-parse-bisent', local_dir_use_symlinks=False)

print('All models downloaded!')
" 2>&1 || warn "Some models may need to be downloaded on first run"

success "MuseTalk models ready"

# ============================================================
# 6. CREATE DIRECTORIES
# ============================================================
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "[6/6] Setting up directories..."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

mkdir -p "$BASE_DIR/server/avatar_videos"
mkdir -p "$BASE_DIR/server/avatar_cache"
mkdir -p "$BASE_DIR/results/server"
mkdir -p "$BASE_DIR/cache"

success "Directories created"

# ============================================================
# INSTALLATION COMPLETE
# ============================================================
echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          INSTALLATION COMPLETE!                              ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""
echo "Installed components:"
echo "  ✓ vLLM + Gemma 3 1B (LLM, ~100ms)"
echo "  ✓ Distil-Whisper PT-BR (STT, ~47ms)"
echo "  ✓ MuseTalk models (lip-sync generation)"
echo "  ✓ Python dependencies"
echo ""
echo "To start the server, run:"
echo ""
echo "    bash server/start.sh"
echo ""
echo "This will:"
echo "  1. Start vLLM with Gemma 3 1B + 5 warmup inferences"
echo "  2. Start Distil-Whisper PT-BR + 5 warmup inferences"
echo "  3. Start MuseTalk server + 5 warmup inferences"
echo "  4. Keep all models in GPU memory for fast inference"
echo ""