#!/bin/bash
# Update a GPU-server clone of agenx-lora-training and run the training.
#
# Usage:
#   bash deploy-and-train.sh
#
# Optional overrides:
#   LORA_DEPLOY_REPO_URL=https://gitea.cyaren.com/cmedina/agenx-lora-training.git
#   LORA_DEPLOY_REPO_DIR=/opt/loras/agenx-lora-training
#   LORA_DEPLOY_BRANCH=main
#

set -euo pipefail

# Load credentials from env file if it exists (override with env vars).
if [ -f /home/cyaren/.deploy.env ]; then
    set -a; source /home/cyaren/.deploy.env; set +a
fi

DEPLOY_USER="${LORA_DEPLOY_USER:-}"
DEPLOY_TOKEN="${LORA_DEPLOY_TOKEN:-}"

# Build the repo URL; embed user:token for HTTPS auth when provided.
# Use local Gitea if available, otherwise fall back to remote
if ping -c 1 -W 1 192.168.50.232 &>/dev/null; then
    REPO_URL="${LORA_DEPLOY_REPO_URL:-http://192.168.50.232:8091/cmedina/agenx-lora-training.git}"
else
    if [ -n "${DEPLOY_USER}" ] && [ -n "${DEPLOY_TOKEN}" ]; then
        REPO_URL="${LORA_DEPLOY_REPO_URL:-https://${DEPLOY_USER}:${DEPLOY_TOKEN}@gitea.cyaren.com/cmedina/agenx-lora-training.git}"
    else
        REPO_URL="${LORA_DEPLOY_REPO_URL:-https://gitea.cyaren.com/cmedina/agenx-lora-training.git}"
    fi
fi
REPO_DIR="${LORA_DEPLOY_REPO_DIR:-/opt/loras/agenx-lora-training}"
BRANCH="${LORA_DEPLOY_BRANCH:-main}"
REMOTE="${LORA_DEPLOY_REMOTE:-origin}"

RUN_USER="${SUDO_USER:-$(id -un)}"
RUN_GROUP="$(id -gn "${RUN_USER}" 2>/dev/null || id -gn)"
REPO_PARENT="$(dirname "${REPO_DIR}")"

echo "=== LoRA Training Git Deploy ==="
echo "Repo:   ${REPO_URL}"
echo "Dir:    ${REPO_DIR}"
echo "Branch: ${BRANCH}"

if [ ! -d "${REPO_PARENT}" ]; then
    sudo install -d -m 0755 -o "${RUN_USER}" -g "${RUN_GROUP}" "${REPO_PARENT}"
fi

if [ ! -d "${REPO_DIR}/.git" ]; then
    git clone --branch "${BRANCH}" "${REPO_URL}" "${REPO_DIR}"
fi

cd "${REPO_DIR}"

# Ensure the remote URL matches the authenticated HTTPS URL.
git remote set-url "${REMOTE}" "${REPO_URL}"

git fetch "${REMOTE}" "${BRANCH}"
git checkout "${BRANCH}"
git pull --ff-only "${REMOTE}" "${BRANCH}"

# Setup Python environment and train
echo ""
echo "=== Setting up Python environment ==="
python3 -m venv venv
source venv/bin/activate
pip install --upgrade pip
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install transformers datasets trl peft accelerate bitsandbytes

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

# Train
echo ""
echo "=== Starting LoRA training ==="
python3 training/scripts/train.py --config training/configs/llama2-7b-lora.yaml

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