#!/bin/bash
# ═══════════════════════════════════════════════════════════════
# Claudia Persistent Absorber v2 — Instance Setup
# Target: Vast.ai 33093662 (A100 80GB, Sweden)
# SSH: ssh -p 13662 root@ssh1.vast.ai
# ═══════════════════════════════════════════════════════════════

set -e

echo "═══════════════════════════════════════════"
echo " Claudia Persistent Absorber v2 — Setup"
echo "═══════════════════════════════════════════"

# ── 1. System packages ──
echo "[1/6] System packages..."
apt-get update -qq && apt-get install -y -qq git wget > /dev/null 2>&1

# ── 2. Python deps ──
echo "[2/6] Python dependencies..."
pip install -q torch torchvision --index-url https://download.pytorch.org/whl/cu124
pip install -q 'transformers>=4.57,<5' accelerate peft safetensors huggingface_hub datasets trl

# ── 3. Create directories ──
echo "[3/6] Creating directories..."
mkdir -p /workspace/models
mkdir -p /workspace/checkpoints
mkdir -p /workspace/logs
mkdir -p /workspace/scripts

# ── 4. Download base model (if not cached) ──
MODEL_DIR="/workspace/models/qwen3-base"
if [ -d "$MODEL_DIR" ] && [ -f "$MODEL_DIR/config.json" ]; then
    echo "[4/6] Base model already cached at $MODEL_DIR"
else
    echo "[4/6] Downloading base model (this takes ~20 min)..."
    python3 -c "
from huggingface_hub import snapshot_download
snapshot_download(
    'huihui-ai/Huihui-Qwen3-Omni-30B-A3B-Instruct-abliterated',
    local_dir='$MODEL_DIR',
    max_workers=4,
)
print('Base model downloaded.')
"
fi

# ── 5. Download Claudia v6 adapter (if not cached) ──
ADAPTER_DIR="/workspace/models/claudia-v6"
if [ -d "$ADAPTER_DIR" ] && [ -f "$ADAPTER_DIR/adapter_config.json" ]; then
    echo "[5/6] Claudia v6 adapter already cached at $ADAPTER_DIR"
else
    echo "[5/6] Downloading Claudia v6 adapter..."
    python3 -c "
from huggingface_hub import snapshot_download
snapshot_download(
    'msrcam/claudia-v6-combined',
    local_dir='$ADAPTER_DIR',
    max_workers=4,
)
print('Adapter downloaded.')
"
fi

# ── 6. Copy absorber script ──
echo "[6/6] Setup complete."
echo ""
echo "═══════════════════════════════════════════"
echo " FIRST RUN:"
echo "   python /workspace/scripts/persistent_absorber.py \\"
echo "     --model_path /workspace/models/qwen3-base \\"
echo "     --adapter_path /workspace/models/claudia-v6 \\"
echo "     --ffn_patch /workspace/models/claudia-v6/ffn_patch.pt"
echo ""
echo " RESUME (after first session):"
echo "   python /workspace/scripts/persistent_absorber.py \\"
echo "     --model_path /workspace/checkpoints/claudia_session_YYYYMMDD_HHMM \\"
echo "     --checkpoint /workspace/checkpoints/claudia_session_YYYYMMDD_HHMM"
echo "═══════════════════════════════════════════"