#!/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 ==="
ai none
sleep 2

# 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..."
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

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

# 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)}')
    print(f'  ✓ VRAM: {torch.cuda.get_device_properties(0).total_mem / 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 "Dataset not found. Downloading..."
    curl -L -o dataset/combined_20k.jsonl https://gitea.cyaren.com/cmedina/agenx-lora-training/raw/branch/main/dataset/combined_20k.jsonl
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 ""

python3 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"