#!/bin/bash

echo "🚀 Complete Base LLM Setup"
echo "=========================="

# Check if virtual environment exists
if [ ! -d "venv" ]; then
    echo "📦 Virtual environment not found. Creating..."
    chmod +x setup.sh
    ./setup.sh
else
    echo "✅ Virtual environment found"
fi

# Activate virtual environment
echo "🔧 Activating virtual environment..."
source venv/bin/activate

# Check HuggingFace token
if [ -z "$HUGGINGFACE_TOKEN" ]; then
    echo "⚠️  HUGGINGFACE_TOKEN not set"
    echo "Please set your token:"
    echo "export HUGGINGFACE_TOKEN='your_token_here'"
    echo ""
    read -p "Enter your HuggingFace token: " token
    if [ ! -z "$token" ]; then
        export HUGGINGFACE_TOKEN="$token"
        echo "✅ Token set"
    else
        echo "❌ No token provided. Exiting..."
        exit 1
    fi
else
    echo "✅ HuggingFace token found"
fi

# Check if model exists
if [ ! -d "models/llama-3.1-8b-instruct" ]; then
    echo "📥 Downloading base model..."
    python scripts/download_model.py
else
    echo "✅ Base model found"
fi

# Check if dataset exists
if [ ! -f "data/training_data.jsonl" ]; then
    echo "📊 Creating sample dataset..."
    python scripts/create_sample_dataset.py
    # Choose option 1 (create sample dataset)
    echo "1" | python scripts/create_sample_dataset.py
else
    echo "✅ Training dataset found"
fi

# Check if config exists
if [ ! -f "configs/llama_config.yaml" ]; then
    echo "⚙️  Creating model configuration..."
    python scripts/download_model.py
else
    echo "✅ Model configuration found"
fi

echo ""
echo "🎉 Setup Complete!"
echo "=================="
echo ""
echo "📋 Next steps:"
echo "1. Review configuration: cat configs/llama_config.yaml"
echo "2. Start fine-tuning: python scripts/finetune_lora.py"
echo "3. Test model: python scripts/test_model.py"
echo "4. Start vLLM service: docker-compose up -d vllm"
echo ""
echo "💡 Tips:"
echo "- Always activate venv: source venv/bin/activate"
echo "- Monitor GPU usage: nvidia-smi"
echo "- Check logs: tail -f logs/training.log"
echo ""
echo "🚀 Ready to start fine-tuning!"