#!/bin/bash
# Train LoRA on this server (assumes GPU is available).
# Run this after deploy-and-train.sh has cloned the repo.

set -e

REPO_DIR="${1:-$(pwd)}"

echo "=== LoRA Training Setup ==="
echo "Repo: ${REPO_DIR}"
echo ""

# Unload AI model to free VRAM
echo "=== Unloading AI model ==="
if command -v ai &> /dev/null; then
    ai none
    sleep 2
else
    echo "Warning: 'ai' command not found, skipping AI unload"
fi

# Check Python
if ! command -v python3 &> /dev/null; then
    echo "ERROR: python3 not found"
    exit 1
fi

echo "[1/4] Setup Python environment..."
python3 -m venv venv
source venv/bin/activate
pip install --upgrade pip

echo "[2/4] Install PyTorch with CUDA (RTX 5090 compatible)..."
pip install --force-reinstall torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

echo "[3/4] Install training dependencies..."
pip install transformers datasets trl peft accelerate bitsandbytes deepspeed

# Check GPU
echo ""
echo "[4/4] Checking GPU..."
python3 -c "
import torch
if torch.cuda.is_available():
    print(f'  ✓ GPU: {torch.cuda.get_device_name(0)}')
    props = torch.cuda.get_device_properties(0)
    print(f'  ✓ VRAM: {props.total_memory / 1e9:.1f} GB')
else:
    print('  ✗ No GPU detected')
    print('  Please ensure CUDA drivers are installed')
    exit(1)
"

echo ""
echo "=== Dataset Preparation ==="
echo "Checking for combined_20k.jsonl..."
if [ ! -f "dataset/combined_20k.jsonl" ]; then
    echo "ERROR: Dataset not found!"
    echo "The file should be in the repo at: dataset/combined_20k.jsonl"
    exit 1
fi

echo "Preparing dataset..."
python3 training/scripts/prepare_dataset.py --input dataset/combined_20k.jsonl --output training/data

echo ""
echo "=== Starting Training ==="
echo "Config: training/configs/llama2-7b-lora.yaml"
echo "Output: training/output/llama2-7b-lora/"
echo ""
echo "Training will take 6-24 hours depending on GPU."
echo "Press Ctrl+C to stop (model will be saved at checkpoint)."
echo ""

# Use torchrun for distributed training (2 GPUs)
torchrun --nproc_per_node=2 training/scripts/train.py --config training/configs/llama2-7b-lora.yaml

echo ""
echo "=== Training Complete ==="
echo "Model saved to: training/output/llama2-7b-lora/"
echo ""
echo "To generate summaries:"
echo "  python3 training/scripts/inference.py \\"
echo "      --model training/output/llama2-7b-lora \\"
echo "      --task \"Your task here\" \\"
echo "      --files src/file.py"