#!/usr/bin/env bash
# scripts/00_env_setup.sh (CUDA 12.1 variant)
# One-shot setup for Ankahi training environment on H100.
# Uses cu121 wheels for better compatibility with newer torchao/unsloth.
# Usage: bash scripts/00_env_setup.sh
set -euo pipefail

echo "============================================================"
echo "  ANKAHI — Environment Setup (CUDA 12.1)"
echo "============================================================"

# ── 1. System checks ─────────────────────────────────────────────
echo "[1/8] Checking system..."
python_version=$(python3 --version 2>&1)
echo "  Python: $python_version"

nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader || {
    echo "ERROR: nvidia-smi failed — is CUDA available?"
    exit 1
}

nvcc_version=$(nvcc --version 2>&1 | grep "release" | awk '{print $5}' | tr -d ',')
echo "  CUDA: $nvcc_version"

# ── 2. Upgrade pip ───────────────────────────────────────────────
echo "[2/8] Upgrading pip..."
python3 -m pip install --upgrade pip setuptools wheel --quiet

# ── 3. PyTorch (CUDA 12.4 wheels) ───────────────────────────────
echo "[3/8] Installing PyTorch 2.6.0 with CUDA 12.4..."
pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 \
    --index-url https://download.pytorch.org/whl/cu124 --quiet

# Verify
python3 -c "import torch; assert torch.cuda.is_available(), 'CUDA not available!'; print(f'  torch {torch.__version__}, CUDA {torch.version.cuda}, device: {torch.cuda.get_device_name(0)}')"

# ── 4. Flash-Attention ────────────────────────────────────────────
echo "[4/8] Installing flash-attn..."
pip install flash-attn>=2.6.0 --no-build-isolation --quiet

# ── 5. Core HuggingFace stack ─────────────────────────────────────
echo "[5/8] Installing HuggingFace stack..."
pip install \
    "transformers==4.57.1" \
    "torchao==0.9.0" \
    "datasets>=3.4.1,<4.4.0" \
    "peft>=0.12.0" \
    "trl>=0.18.2,<=0.24.0" \
    "accelerate>=0.30.0" \
    "bitsandbytes>=0.43.0" \
    wandb \
    --quiet

# ── 6. Unsloth ────────────────────────────────────────────────────
echo "[6/8] Installing Unsloth (Gemma 4 support)..."
pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" 

python3 -c "from unsloth import FastVisionModel; print('  Unsloth FastVisionModel OK')"

# ── 7. Audio, TTS, and data deps ─────────────────────────────────
echo "[7/8] Installing audio/TTS/data dependencies..."
# numpy pinned to 1.26.4 for f5-tts compatibility
pip install \
    "numpy<=1.26.4" \
    librosa soundfile ffmpeg-python audiomentations \
    TTS \
    sacrebleu \
    "sentence-transformers>=3.0.0" \
    pandas Pillow pyarrow requests tqdm pyyaml python-dotenv \
    anthropic "google-generativeai>=0.7.0" \
    pytest pytest-cov black ruff \
    --quiet

# AI4Bharat Indic-TTS (clone if not present)
if [ ! -d "vendor/Indic-TTS" ]; then
    echo "  Cloning AI4Bharat Indic-TTS..."
    mkdir -p vendor
    git clone --depth 1 https://github.com/AI4Bharat/Indic-TTS vendor/Indic-TTS --quiet
    pip install -e vendor/Indic-TTS --quiet || echo "  (Indic-TTS install needs manual steps — see vendor/Indic-TTS/README)"
fi

# ── 8. Install ankahi package itself (editable) ───────────────────
echo "[8/8] Installing ankahi in editable mode..."
pip install -e . --quiet

echo ""
echo "============================================================"
echo "  Setup complete. Run sanity check next:"
echo "  python scripts/07_overfit_sanity_check.py"
echo "============================================================"