#!/usr/bin/env bash
# =============================================================================
#  LUNA 300M — Cloud Setup & Train Entrypoint (One Command)
#  Downloads dataset, pretrains 300M model, validates every 100 steps,
#  uploads all checkpoints to HuggingFace at the end.
#
#  USAGE:
#    bash setup_and_train_300m.sh huggingface ASTERIZER/Luna_Dataset
#    bash setup_and_train_300m.sh huggingface ASTERIZER/Luna_Dataset 10000000  # smoke test
#    bash setup_and_train_300m.sh local /workspace/data/litdata_pretrain_final
#
#  REQUIRED ENV:
#    HF_TOKEN=hf_xxx   (for uploading checkpoints to HuggingFace)
# =============================================================================

set -e

DATA_SOURCE="${1:-huggingface}"
DATA_ID="${2:-ASTERIZER/Luna_Dataset}"
MAX_TOKENS="${3:-4515286950}"
DATA_DIR="/workspace/data/litdata_pretrain_final"
OUT_DIR="/workspace/out/pretrain/luna-300m"

echo "=========================================="
echo "  LUNA 300M — Cloud Setup & Pretrain"
echo "  Source  : $DATA_SOURCE"
echo "  ID/Path : $DATA_ID"
echo "  Tokens  : $MAX_TOKENS"
echo "  Output  : $OUT_DIR"
echo "=========================================="

# ── 1. Python packages ────────────────────────────────────────────────────────
echo ""
echo "[1/5] Installing dependencies..."

pip install -q --upgrade pip
pip install -q \
    torch torchvision \
    psutil \
    pyyaml \
    huggingface_hub \
    gdown \
    tensorboard \
    transformers \
    datasets \
    pyarrow 2>/dev/null || true

echo "      Done."

# ── 2. Download dataset ───────────────────────────────────────────────────────
echo ""
echo "[2/5] Fetching dataset..."

if [ "$DATA_SOURCE" = "gdrive" ]; then
    python fetch_data.py --source gdrive --gdrive_id "$DATA_ID" --out_dir "$DATA_DIR"
elif [ "$DATA_SOURCE" = "huggingface" ]; then
    HF_TOKEN="${HF_TOKEN:-}"
    python fetch_data.py --source huggingface --hf_repo "$DATA_ID" --out_dir "$DATA_DIR" --hf_token "$HF_TOKEN"
elif [ "$DATA_SOURCE" = "local" ]; then
    python fetch_data.py --source local --local_path "$DATA_ID" --out_dir "$DATA_DIR"
else
    echo "Unknown source: $DATA_SOURCE (use: gdrive | huggingface | local)"
    exit 1
fi

DATA_DIR_ENV="$DATA_DIR" python - <<'PYEOF'
from pathlib import Path
import json
import os

data_dir = Path(os.environ["DATA_DIR_ENV"])
index_path = data_dir / "index.json"
if not index_path.exists():
    candidates = sorted(data_dir.glob("**/index.json"))
    if candidates:
        print("ERROR: index.json still not at dataset root after fetch.")
        print("Found candidates:")
        for candidate in candidates[:10]:
            print(f"  - {candidate}")
    else:
        print(f"ERROR: No index.json found under {data_dir}")
    raise SystemExit(1)

with open(index_path, encoding="utf-8") as f:
    idx = json.load(f)

chunks = idx.get("chunks", [])
total = sum(c.get("dim", 0) for c in chunks)
print(f"  Dataset ready: {len(chunks)} chunks, {total:,} tokens at {index_path}")
PYEOF

# ── 3. System probe ──────────────────────────────────────────────────────────
echo ""
echo "[3/5] System probe..."
python - <<'PYEOF'
import os
import psutil
import torch

props = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None
if props:
    print(f"  GPU  : {props.name} ({props.total_memory / 1024**3:.1f} GB)")
else:
    print("  GPU: None")
print(f"  RAM  : {psutil.virtual_memory().total / 1024**3:.1f} GB")
print(f"  CPUs : {os.cpu_count()}")
PYEOF

# ── 4. Pretrain (with validation every 100 steps) ─────────────────────────────
echo ""
echo "[4/5] Starting LUNA 300M pretraining..."
echo "      - Validation: every 100 steps (10 generation prompts)"
echo "      - Checkpoints: every 1000 steps"
echo "      - Auto-detects VRAM and fills to maximum utilization"
echo ""

python train_300m.py \
    --config train_config_300m.yaml \
    --data_path "$DATA_DIR" \
    --out_dir "$OUT_DIR" \
    --max_tokens "$MAX_TOKENS"

# ── 5. Summary ────────────────────────────────────────────────────────────────
echo ""
echo "=========================================="
echo "  LUNA 300M — COMPLETE"
echo "  Output      : $OUT_DIR"
echo "  Final model : $OUT_DIR/final/lit_model.pth"
if [ -n "$HF_TOKEN" ]; then
echo "  HF Upload   : https://huggingface.co/ASTERIZER/LUNA-300M"
else
echo "  HF Upload   : SKIPPED (set HF_TOKEN to enable)"
fi
echo "=========================================="