#!/bin/bash
#
# dumont-talker - Installation Script
# Installs all dependencies for the speech-to-speech avatar system
#

set -e

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

GREEN='\033[0;32m'
CYAN='\033[0;36m'
NC='\033[0m'

success() { echo -e "${GREEN}✓${NC} $1"; }
header() { echo -e "\n${CYAN}$1${NC}"; }

echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          DUMONT-TALKER - INSTALLATION                        ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""

# ============================================================
# 1. PYTORCH WITH BLACKWELL SUPPORT
# ============================================================
header "[1/4] Installing PyTorch with Blackwell GPU support..."

PYTORCH_OK=$(python3 -c "
import torch
if torch.cuda.is_available():
    arch = torch.cuda.get_arch_list()
    if 'sm_120' in arch or 'sm_100' in arch:
        print('ok')
" 2>/dev/null || echo "")

if [ "$PYTORCH_OK" = "ok" ]; then
    success "PyTorch already has Blackwell support"
else
    pip install --force-reinstall torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128 --quiet
    success "PyTorch installed with CUDA 12.8"
fi

# ============================================================
# 2. SERVER DEPENDENCIES
# ============================================================
header "[2/4] Installing server dependencies..."

pip install --quiet \
    fastapi uvicorn httpx python-multipart pydub aiofiles \
    av aiortc websockets \
    diffusers accelerate transformers huggingface_hub \
    opencv-python-headless soundfile omegaconf ffmpeg-python

success "Server dependencies installed"

# ============================================================
# 3. TESTING TOOLS
# ============================================================
header "[3/4] Installing testing tools..."

pip install playwright --quiet
playwright install chromium --quiet 2>/dev/null || playwright install chromium
playwright install-deps chromium --quiet 2>/dev/null || true

success "Playwright installed"

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

check_model() {
    if [ -L "$1" ] && [ -e "$1" ]; then
        echo -e "  ${GREEN}✓${NC} $2"
    elif [ -d "$1" ]; then
        echo -e "  ${GREEN}✓${NC} $2"
    else
        echo -e "  ✗ $2 (not found)"
    fi
}

check_model "$PROJECT_DIR/models/avatar/musetalk" "MuseTalk UNet"
check_model "$PROJECT_DIR/models/avatar/vae" "VAE"
check_model "$PROJECT_DIR/models/llm/gemma" "Gemma 3 1B"
check_model "$PROJECT_DIR/models/stt/whisper" "Whisper"

# ============================================================
# COMPLETE
# ============================================================
echo ""
echo "╔══════════════════════════════════════════════════════════════╗"
echo "║          INSTALLATION COMPLETE!                              ║"
echo "╚══════════════════════════════════════════════════════════════╝"
echo ""
echo "  To start: $SCRIPT_DIR/start.sh"
echo ""