# Speech-X — First-time environment setup (PowerShell)
# Creates the conda 'avatar' environment and installs all dependencies in stages.
# Run from the repo root: .\setup\setup.ps1

$ErrorActionPreference = "Stop"
$CONDA_ENV = "avatar"

Write-Host "=== Speech-X Setup (conda env: $CONDA_ENV) ===" -ForegroundColor Cyan

# ── Stage 0: sanity checks ───────────────────────────────────────────────────
if (-not (Get-Command conda -ErrorAction SilentlyContinue)) {
    Write-Error "conda not found. Install Miniconda: https://docs.conda.io/en/latest/miniconda.html"
    exit 1
}

if (Get-Command nvidia-smi -ErrorAction SilentlyContinue) {
    Write-Host "GPU detected:"
    nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
} else {
    Write-Warning "nvidia-smi not found — CPU-only mode."
}

# ── Stage 1: Create conda environment ────────────────────────────────────────
$envList = conda env list
if ($envList -match "(?m)^$CONDA_ENV\s") {
    Write-Host "Conda env '$CONDA_ENV' already exists — skipping creation."
} else {
    Write-Host "Creating conda env '$CONDA_ENV' (Python 3.12)..."
    conda create -y -n $CONDA_ENV python=3.12
}

# ── Stage 2: PyTorch ─────────────────────────────────────────────────────────
Write-Host "Installing PyTorch 2.5.1 + CUDA 12.4..."
conda run -n $CONDA_ENV pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 `
    --index-url https://download.pytorch.org/whl/cu124

# ── Stage 3: MMLab packages ──────────────────────────────────────────────────
Write-Host "Installing MMLab packages..."
conda run -n $CONDA_ENV pip install --no-cache-dir -U openmim
conda run -n $CONDA_ENV mim install mmengine
try {
    conda run -n $CONDA_ENV mim install "mmcv==2.2.0"
} catch {
    conda run -n $CONDA_ENV pip install "mmcv-lite==2.2.0"
}
conda run -n $CONDA_ENV mim install "mmdet==3.3.0"
conda run -n $CONDA_ENV mim install "mmpose==1.3.0"

# ── Stage 4: MuseTalk core deps ──────────────────────────────────────────────
Write-Host "Installing MuseTalk dependencies..."
conda run -n $CONDA_ENV pip install `
    diffusers==0.30.2 `
    accelerate==0.28.0 `
    "numpy==1.26.4" `
    "opencv-python==4.10.0.84" `
    "soundfile==0.12.1" `
    "transformers==4.39.2" `
    "huggingface_hub==0.30.2" `
    "librosa==0.10.2" `
    "einops==0.8.1" `
    gdown requests `
    "imageio==2.34.0" imageio-ffmpeg `
    "omegaconf==2.3.0" ffmpeg-python moviepy

# ── Stage 5: Project-specific deps ───────────────────────────────────────────
Write-Host "Installing project dependencies..."
conda run -n $CONDA_ENV pip install -r backend/requirements.txt

# ── Frontend ─────────────────────────────────────────────────────────────────
Write-Host "Installing frontend dependencies..."
Push-Location frontend
npm install
Pop-Location

Write-Host ""
Write-Host "=== Setup complete ===" -ForegroundColor Green
Write-Host ""
Write-Host "Activate:       conda activate $CONDA_ENV"
Write-Host "Avatar page:    conda activate $CONDA_ENV; cd backend; python api/server.py"
Write-Host "Voice agent:    conda activate $CONDA_ENV; cd backend; python agent.py dev"
Write-Host "Frontend:       cd frontend; npm run dev"