#!/bin/bash
# ============================================================================
# Daimon Training — RunPod First-Boot Setup
# Liberation Labs
# ============================================================================
#
# Run this ONCE on a fresh RunPod pod before training.
# It installs dependencies, pulls the model & data, and verifies the environment.
#
# Requirements:
#   - 1x H200 SXM 141GB
#   - 188GB+ system RAM (critical for CPU-offloaded optimizer)
#   - 400GB+ disk on /workspace (persistent volume)
#   - HF_TOKEN environment variable set (model is gated)
#
# Usage:
#   export HF_TOKEN="hf_your_token_here"
#   bash /workspace/runpod-template/setup.sh
# ============================================================================

set -e

echo "============================================================"
echo "  DAIMON FULL-PARAMETER SFT — POD SETUP"
echo "  Liberation Labs"
echo "  $(date)"
echo "============================================================"

# ── 1. Find Python ──────────────────────────────────────────────────────────
export PATH=/opt/conda/bin:/usr/local/bin:$PATH
PYTHON=$(which python3.11 2>/dev/null || which python3 2>/dev/null)
echo "Python: $PYTHON ($($PYTHON --version 2>&1))"

# ── 2. Check GPU ──────────────────────────────────────────────────────────
echo ""
echo "=== GPU Check ==="
GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
echo "GPUs detected: $GPU_COUNT"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader

if [ "$GPU_COUNT" -lt 1 ]; then
    echo ""
    echo "FATAL: No GPUs detected."
    exit 1
fi

# Check VRAM (need >= 140GB for full SFT on single GPU)
VRAM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1 | tr -d ' ')
echo "VRAM: ${VRAM} MiB"
if [ "$VRAM" -lt 140000 ]; then
    echo ""
    echo "FATAL: GPU has ${VRAM} MiB VRAM. Need >= 140,000 MiB."
    echo ""
    echo "Why: Full-parameter SFT memory budget:"
    echo "  Model params (bf16):    ~70GB → GPU"
    echo "  Activations (grad ckpt): ~20GB → GPU"
    echo "  Total GPU:              ~90GB of 141GB"
    echo ""
    echo "Fix: Provision a pod with 1x H200 SXM 141GB."
    exit 1
fi

# ── 3. Check system RAM ────────────────────────────────────────────────────
echo ""
echo "=== System RAM ==="
TOTAL_RAM=$(free -g | grep Mem | awk '{print $2}')
echo "Total: ${TOTAL_RAM} GB"

# System RAM is CRITICAL for full SFT — optimizer states and gradients are CPU-offloaded
if [ "$TOTAL_RAM" -lt 180 ]; then
    echo "FATAL: System RAM is ${TOTAL_RAM}GB. Need >= 180GB."
    echo ""
    echo "Why: Full SFT CPU-offloaded memory budget:"
    echo "  Gradients (bf16):        ~70GB → CPU"
    echo "  Adafactor optimizer:     ~35GB → CPU"
    echo "  Total CPU:              ~105GB"
    echo "  Plus OS/data overhead:   ~30GB"
    echo ""
    echo "AdamW is NOT viable — its fp32 states would need ~280GB CPU RAM."
    echo "Even Adafactor needs ~105GB + headroom."
    echo ""
    echo "Fix: Provision a pod with >= 188GB system RAM."
    exit 1
fi

# ── 4. Check disk space ────────────────────────────────────────────────────
echo ""
echo "=== Disk Space ==="
AVAIL_GB=$(df -BG /workspace | tail -1 | awk '{print $4}' | tr -d 'G')
echo "Available on /workspace: ${AVAIL_GB} GB"

if [ "$AVAIL_GB" -lt 300 ]; then
    echo "FATAL: Less than 300GB on /workspace."
    echo "Full SFT needs: model (~70GB) + data + checkpoints (~210GB for 3 × 70GB)."
    echo "Mount a 400GB+ persistent volume."
    exit 1
elif [ "$AVAIL_GB" -lt 400 ]; then
    echo "WARNING: Less than 400GB on /workspace."
    echo "Full model checkpoints are ~70GB each. With save_total_limit=3, need ~210GB."
    echo "Will be tight — consider a larger volume."
fi

# ── 5. Check HF_TOKEN ──────────────────────────────────────────────────────
echo ""
echo "=== HuggingFace Authentication ==="
if [ -z "$HF_TOKEN" ]; then
    echo "FATAL: HF_TOKEN environment variable not set."
    echo "Qwen3.6-35B-A3B is a gated model. You need a HuggingFace token."
    echo ""
    echo "Fix: export HF_TOKEN='hf_your_token_here'"
    echo "Or set it in the RunPod pod template environment variables."
    exit 1
else
    echo "HF_TOKEN is set (${#HF_TOKEN} chars)"
fi

# ── 6. Install dependencies (pinned versions) ─────────────────────────────
echo ""
echo "=== Installing Dependencies ==="
$PYTHON -m pip install --upgrade -q pip

echo "Installing PyTorch..."
$PYTHON -m pip install -q \
    torch==2.7.1 \
    torchvision==0.22.1 \
    --index-url https://download.pytorch.org/whl/cu124 \
    2>&1 | tail -2

echo "Installing training stack (pinned versions)..."
$PYTHON -m pip install -q \
    transformers==5.12.1 \
    trl==1.7.0 \
    datasets==5.0.0 \
    accelerate==1.14.0 \
    deepspeed==0.16.7 \
    safetensors==0.8.0 \
    pyyaml==6.0.2 \
    2>&1 | tail -3

echo "Installing flash-attn (may take a few minutes to compile)..."
$PYTHON -m pip install -q flash-attn --no-build-isolation 2>&1 | tail -3 || {
    echo "WARNING: flash-attn failed to install. Will fall back to SDPA attention."
    echo "This is fine — SDPA is only ~5% slower on H200."
}

echo "Dependencies installed."

# ── 7. Verify critical packages ────────────────────────────────────────────
echo ""
echo "=== Package Verification ==="
$PYTHON -c "
import torch, transformers, trl, datasets, accelerate, deepspeed, safetensors
print(f'torch:        {torch.__version__}')
print(f'transformers: {transformers.__version__}')
print(f'trl:          {trl.__version__}')
print(f'datasets:     {datasets.__version__}')
print(f'accelerate:   {accelerate.__version__}')
print(f'deepspeed:    {deepspeed.__version__}')
print(f'safetensors:  {safetensors.__version__}')
print(f'CUDA:         {torch.version.cuda}')
print(f'GPUs:         {torch.cuda.device_count()}')
try:
    import flash_attn
    print(f'flash_attn:   {flash_attn.__version__}')
except ImportError:
    print('flash_attn:   not installed (using SDPA fallback)')
"

# ── 8. Verify Qwen3.6 architecture support ─────────────────────────────────
echo ""
echo "=== Model Architecture Check ==="
MODEL_REVISION="995ad96eacd98c81ed38be0c5b274b04031597b0"
$PYTHON -c "
from transformers import AutoConfig
c = AutoConfig.from_pretrained('Qwen/Qwen3.6-35B-A3B', revision='$MODEL_REVISION', trust_remote_code=True)
print(f'Model type:      {c.model_type}')
print(f'Hidden size:     {c.hidden_size}')
print(f'Num layers:      {c.num_hidden_layers}')
print(f'Num experts:     {getattr(c, \"num_experts\", \"N/A\")}')
print(f'Vocab size:      {c.vocab_size}')
print(f'Pinned revision: $MODEL_REVISION')
print('Architecture supported: OK')
" || {
    echo "FATAL: Qwen3.6 architecture not supported by installed transformers."
    echo "Upgrade: pip install --upgrade transformers"
    exit 1
}

# ── 9. Pull model from HuggingFace ─────────────────────────────────────────
echo ""
echo "=== Model Download ==="
MODEL_DIR="/workspace/models/Qwen3.6-35B-A3B"
if [ -d "$MODEL_DIR" ] && [ -f "$MODEL_DIR/config.json" ]; then
    echo "Model already downloaded at $MODEL_DIR"
else
    echo "Downloading Qwen3.6-35B-A3B (~70GB, this will take a while)..."
    mkdir -p /workspace/models
    $PYTHON -c "
from huggingface_hub import snapshot_download
import os
snapshot_download(
    'Qwen/Qwen3.6-35B-A3B',
    revision='$MODEL_REVISION',
    local_dir='$MODEL_DIR',
    token=os.environ['HF_TOKEN'],
)
print('Model download complete.')
"
fi

# ── 10. Pull training data ──────────────────────────────────────────────────
echo ""
echo "=== Training Data ==="
DATA_DIR="/workspace/daimon-data"
mkdir -p "$DATA_DIR"

if [ -d "$DATA_DIR/train_arrow" ] && [ -d "$DATA_DIR/valid_arrow" ]; then
    echo "Arrow data already present. Verifying..."
    $PYTHON -c "
from datasets import load_from_disk
t = load_from_disk('$DATA_DIR/train_arrow')
v = load_from_disk('$DATA_DIR/valid_arrow')
print(f'Train: {len(t):,} samples | Valid: {len(v):,} samples — OK')
"
else
    echo "Downloading and preparing training data..."
    $PYTHON -c "
import os, json, gzip, shutil
from huggingface_hub import hf_hub_download, list_repo_files
from datasets import Dataset, load_dataset

DATA_DIR = '$DATA_DIR'
REPO = 'HumboldtJoker/daimon-sft-data'
token = os.environ.get('HF_TOKEN')

try:
    # Try loading as a HF dataset first
    ds = load_dataset(REPO, token=token)
    if 'train' in ds:
        ds['train'].save_to_disk(f'{DATA_DIR}/train_arrow')
        print(f'Train: {len(ds[\"train\"]):,} samples saved as Arrow')
    if 'validation' in ds:
        ds['validation'].save_to_disk(f'{DATA_DIR}/valid_arrow')
        print(f'Valid: {len(ds[\"validation\"]):,} samples saved as Arrow')
    elif 'test' in ds:
        ds['test'].save_to_disk(f'{DATA_DIR}/valid_arrow')
        print(f'Valid: {len(ds[\"test\"]):,} samples saved as Arrow')
    else:
        # Split train into train/valid
        split = ds['train'].train_test_split(test_size=0.05, seed=42)
        split['train'].save_to_disk(f'{DATA_DIR}/train_arrow')
        split['test'].save_to_disk(f'{DATA_DIR}/valid_arrow')
        print(f'Auto-split: Train {len(split[\"train\"]):,} | Valid {len(split[\"test\"]):,}')
except Exception as e:
    print(f'HF dataset load failed: {e}')
    print('Trying file-based download...')

    # Fall back to downloading individual files
    try:
        files = list_repo_files(REPO, repo_type='dataset', token=token)
        for f in files:
            if f.endswith(('.jsonl', '.jsonl.gz', '.json')):
                print(f'Downloading {f}...')
                hf_hub_download(REPO, f, repo_type='dataset', local_dir=DATA_DIR, token=token)
    except Exception as e2:
        print(f'File download also failed: {e2}')
        print('DATA MUST BE UPLOADED MANUALLY to {DATA_DIR}/')
        print('Expected format: JSONL with {\"messages\": [{\"role\": ..., \"content\": ...}, ...]}')

    # Convert any JSONL files to Arrow
    for split_name in ['train', 'valid']:
        jsonl = f'{DATA_DIR}/{split_name}.jsonl'
        gz = f'{DATA_DIR}/{split_name}.jsonl.gz'
        arrow_dir = f'{DATA_DIR}/{split_name}_arrow'

        if os.path.exists(gz) and not os.path.exists(jsonl):
            with gzip.open(gz, 'rb') as fin, open(jsonl, 'wb') as fout:
                shutil.copyfileobj(fin, fout)

        if os.path.exists(jsonl) and not os.path.exists(arrow_dir):
            data = []
            with open(jsonl) as fh:
                for line in fh:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        d = json.loads(line)
                        if 'messages' in d and len(d['messages']) >= 2:
                            data.append(d)
                    except:
                        pass
            ds = Dataset.from_list(data)
            ds.save_to_disk(arrow_dir)
            print(f'{split_name}: {len(data):,} examples saved as Arrow')

# Final verification
try:
    from datasets import load_from_disk
    t = load_from_disk(f'{DATA_DIR}/train_arrow')
    print(f'Verified train: {len(t):,} samples')
    if os.path.isdir(f'{DATA_DIR}/valid_arrow'):
        v = load_from_disk(f'{DATA_DIR}/valid_arrow')
        print(f'Verified valid: {len(v):,} samples')
except:
    print('WARNING: Could not verify data. Check $DATA_DIR manually.')
"
fi

# ── 11. Create persistent directories ──────────────────────────────────────
echo ""
echo "=== Creating Directories ==="
mkdir -p /workspace/daimon-sft/logs
mkdir -p /workspace/daimon-sft/checkpoints
echo "Output directories created on persistent volume."

# ── 12. Copy template files to /workspace ───────────────────────────────────
echo ""
echo "=== Copying Template Files ==="
SCRIPT_DIR=$(dirname "$(readlink -f "$0")")
cp "$SCRIPT_DIR/train_daimon.py" /workspace/runpod-template/train_daimon.py 2>/dev/null || true
cp "$SCRIPT_DIR/train_daimon_config.yaml" /workspace/runpod-template/train_daimon_config.yaml 2>/dev/null || true
cp "$SCRIPT_DIR/ds_config_zero2.json" /workspace/runpod-template/ds_config_zero2.json 2>/dev/null || true
cp "$SCRIPT_DIR/launch.sh" /workspace/runpod-template/launch.sh 2>/dev/null || true
cp "$SCRIPT_DIR/test_template.py" /workspace/runpod-template/test_template.py 2>/dev/null || true
echo "Template files in /workspace/runpod-template/"

# ── 13. Add SSH key for remote access ──────────────────────────────────────
echo ""
echo "=== SSH Key ==="
mkdir -p ~/.ssh
echo "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIOtjekz8l1s6xTAXlhZJg/A0N3d6mZAyF/EyrEMiBCDG thomas@coalition" >> ~/.ssh/authorized_keys
chmod 700 ~/.ssh
chmod 600 ~/.ssh/authorized_keys
echo "SSH key added."

# ── 14. Clean up HF token from disk cache ──────────────────────────────────
echo ""
echo "=== Security Cleanup ==="
rm -f ~/.cache/huggingface/token 2>/dev/null || true
echo "Cleared cached HF token from disk."

# ── 15. Summary ─────────────────────────────────────────────────────────────
echo ""
echo "============================================================"
echo "  SETUP COMPLETE — FULL-PARAMETER SFT"
echo ""
echo "  Memory budget:"
echo "    GPU:  ~90GB of 141GB (model + activations)"
echo "    CPU: ~105GB of ${TOTAL_RAM}GB (gradients + Adafactor)"
echo ""
echo "  Next steps:"
echo "    1. Run validation:  python3 /workspace/runpod-template/test_template.py"
echo "    2. Start training:  bash /workspace/runpod-template/launch.sh"
echo ""
echo "  Monitor:"
echo "    watch -n 5 nvidia-smi"
echo "    tail -f /workspace/daimon-sft/logs/training_*.log"
echo "============================================================"