#!/bin/bash
#
# LoRA Training Deployment Script
# Clones the repo and trains the LoRA adapter on the server
#
# Usage:
#   bash deploy-and-train.sh
#
# This will:
# 1. Clone the repo to /opt/loras/agenx-lora-training
# 2. Setup Python environment with GPU support
# 3. Prepare the dataset
# 4. Train the LoRA adapter
#

set -e

# Configuration
REPO_URL="https://gitea.cyaren.com/cmedina/agenx-lora-training.git"
INSTALL_DIR="/opt/loras/agenx-lora-training"
PYTHON_VERSION="3.10"

echo "=============================================="
echo "LoRA Training Deployment"
echo "=============================================="
echo ""

# Step 1: Create installation directory
echo "[1/5] Creating installation directory..."
mkdir -p /opt/loras
if [ -d "$INSTALL_DIR" ]; then
    echo "  Directory already exists: $INSTALL_DIR"
    echo "  Removing old clone..."
    rm -rf "$INSTALL_DIR"
fi

# Step 2: Clone the repository
echo "[2/5] Cloning repository..."
git clone "$REPO_URL" "$INSTALL_DIR"
echo "  Repository cloned to: $INSTALL_DIR"

# Step 3: Setup Python environment
echo "[3/5] Setting up Python environment..."
cd "$INSTALL_DIR"

# Check if Python is available
if ! command -v python3 &> /dev/null; then
    echo "ERROR: Python3 not found. Please install Python $PYTHON_VERSION or later."
    exit 1
fi

# Create virtual environment
python3 -m venv venv
source venv/bin/activate

# Upgrade pip
pip install --upgrade pip

# Install PyTorch with CUDA support
echo "  Installing PyTorch with CUDA support..."
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# Install other dependencies
echo "  Installing training dependencies..."
pip install transformers datasets trl peft accelerate bitsandbytes

# Verify GPU availability
echo ""
echo "  Checking GPU availability..."
python3 -c "
import torch
if torch.cuda.is_available():
    print(f'  ✓ CUDA available: {torch.cuda.get_device_name(0)}')
    print(f'  ✓ VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB')
else:
    print('  ✗ CUDA not available. GPU training will not work.')
    print('  Please ensure CUDA drivers are installed.')
    exit(1)
"

# Step 4: Prepare dataset
echo "[4/5] Preparing dataset..."
python3 training/scripts/prepare_dataset.py \
    --input dataset/combined_20k.jsonl \
    --output training/data

echo "  Dataset prepared:"
echo "    - training/data/train.jsonl"
echo "    - training/data/test.jsonl"

# Step 5: Train the model
echo "[5/5] Starting training..."
echo ""
echo "  Training configuration:"
echo "    - Model: meta-llama/Llama-2-7b-hf"
echo "    - Method: QLoRA (4-bit quantization)"
echo "    - Epochs: 3"
echo "    - Batch size: 4"
echo "    - Learning rate: 2e-4"
echo ""
echo "  Estimated training time: 6-24 hours (depending on GPU)"
echo ""

python3 training/scripts/train.py \
    --config training/configs/llama2-7b-lora.yaml

echo ""
echo "=============================================="
echo "Training complete!"
echo "=============================================="
echo ""
echo "Trained model saved to: $INSTALL_DIR/training/output/llama2-7b-lora/"
echo ""
echo "To generate summaries:"
echo "  cd $INSTALL_DIR"
echo "  source venv/bin/activate"
echo "  python3 training/scripts/inference.py \\"
echo "      --model training/output/llama2-7b-lora \\"
echo "      --task \"Your task here\" \\"
echo "      --files src/file.py \\"
echo "      --tests-run --test-count 100"
echo ""