#!/usr/bin/env bash
# =============================================================================
# Indic Heritage Studio v2 — Day 1 environment setup (NVIDIA 8×80GB dev box)
# =============================================================================
# Run from the project root:
#   chmod +x scripts/day1_setup.sh && ./scripts/day1_setup.sh
#
# This script:
#   1. Detects Python version (must be 3.11)
#   2. Creates a venv
#   3. Installs PyTorch + project deps
#   4. Downloads all v2 models (~35 GB to ~/.cache/huggingface)
#   5. Verifies multi-GPU visibility (8 GPUs expected)
#   6. Smoke-tests each pipeline
# =============================================================================
set -e
cd "$(dirname "$0")/.."

echo "=== Indic Heritage Studio v2 — Day 1 Setup ==="
echo "Host: $(hostname) | Date: $(date)"

# --- Python version check ---
PY=$(command -v python3.11 || command -v python3)
if [[ -z "$PY" ]]; then
  echo "ERROR: Python 3.11 not found. Install: sudo apt install python3.11 python3.11-venv"
  exit 1
fi
echo "Python: $($PY --version)"

# --- venv ---
if [[ ! -d .venv ]]; then
  echo "Creating virtualenv…"
  $PY -m venv .venv
fi
source .venv/bin/activate

pip install --upgrade pip wheel

# --- PyTorch (CUDA 12.1) ---
if ! python -c "import torch" 2>/dev/null; then
  echo "Installing PyTorch 2.4.1 + CUDA 12.1…"
  pip install torch==2.4.1 torchvision==0.19.1 --index-url https://download.pytorch.org/whl/cu121
fi

# --- Project deps ---
echo "Installing project dependencies…"
pip install -r requirements.txt

# --- Environment file ---
if [[ ! -f .env ]]; then
  cp .env.example .env
  echo "Created .env from template. Edit it to set AMD_MODEL_API_KEY."
fi

# --- GPU verification ---
echo ""
echo "=== GPU Verification ==="
python -c "
import torch
print(f'PyTorch: {torch.__version__}')
print(f'CUDA available: {torch.cuda.is_available()}')
print(f'GPU count: {torch.cuda.device_count()}')
for i in range(torch.cuda.device_count()):
    print(f'  GPU {i}: {torch.cuda.get_device_name(i)} '
          f'({torch.cuda.get_device_properties(i).total_memory / 1e9:.1f} GB)')
"

# --- Download models ---
echo ""
echo "=== Model Download (~35 GB) ==="
read -p "Download all v2 models now? [y/N] " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
  python scripts/download_models.py
fi

# --- Smoke test ---
echo ""
echo "=== Smoke Test ==="
read -p "Run smoke test (loads SDXL pipeline, generates 1 image)? [y/N] " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
  python -m core.text_to_image \
      --prompt "a peacock in a monsoon garden" \
      --style madhubani \
      --out outputs/smoke_test.png \
      --no-lora --steps 10 --size 512
  echo "Smoke test complete: outputs/smoke_test.png"
fi

echo ""
echo "=== Done. Next steps ==="
echo "  1. Edit .env and set AMD_MODEL_API_KEY"
echo "  2. Place reference heritage art images in assets/styles/<style>_ref.png"
echo "  3. Place raw heritage art images in assets/datasets/raw/<style>/"
echo "  4. Run: python training/prepare_dataset.py"
echo "  5. Run: python -m training.train_lora --style madhubani"
echo "  6. Run: python app.py  # opens UI at http://localhost:7860"