#!/bin/bash
# THE BAKER — Flux.2 LoRA Training Quick Start
# 
# This script walks through the complete training workflow.
# Run from /home/workspaces/thebaker/training/

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

echo "=========================================="
echo "THE BAKER — Flux.2 LoRA Training"
echo "=========================================="
echo ""

# Check environment
echo "1. Checking environment..."
if [ ! -d "/home/workspaces/cultguard-agents/.devenv/state/venv" ]; then
    echo "ERROR: cultguard-agents venv not found"
    echo "Please ensure the devenv is set up first."
    exit 1
fi

# Activate environment
source /home/workspaces/cultguard-agents/.devenv/state/venv/bin/activate

# Check GPU
echo "2. Checking GPU..."
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || {
    echo "WARNING: NVIDIA GPU not detected. Training will fail."
    echo "Proceed anyway? (y/n)"
    read -r response
    if [[ "$response" != "y" ]]; then
        exit 1
    fi
}

# Check dependencies
echo "3. Checking Python dependencies..."
python -c "import torch; import diffusers; import peft; import accelerate" 2>/dev/null || {
    echo "Installing required packages..."
    pip install torch diffusers[torch] transformers accelerate peft pillow
}

echo ""
echo "=========================================="
echo "Choose workflow:"
echo "=========================================="
echo ""
echo "  1) Generate dataset only (create training images)"
echo "  2) Train LoRA (requires generated dataset)"
echo "  3) Generate deck visuals (requires trained LoRA)"
echo "  4) Full pipeline (dataset → train → generate)"
echo "  5) Skip to deck visuals with custom prompt"
echo ""
read -p "Enter choice [1-5]: " choice

case $choice in
    1)
        echo ""
        echo "Generating training dataset..."
        python generate_dataset.py --config config.json
        echo ""
        echo "✓ Dataset complete!"
        echo "  Images: $SCRIPT_DIR/dataset/images/"
        echo "  Captions: $SCRIPT_DIR/dataset/captions/"
        echo ""
        echo "Next step: Run training (option 2)"
        ;;
    
    2)
        echo ""
        echo "Starting LoRA training..."
        echo "This will take 2-4 hours on RTX 6000 Ada."
        echo ""
        read -p "Continue? (y/n): " confirm
        if [[ "$confirm" != "y" ]]; then
            echo "Aborted."
            exit 0
        fi
        
        python train_flux_lora.py --config config.json
        
        echo ""
        echo "✓ Training complete!"
        echo "  LoRA weights: $SCRIPT_DIR/outputs/lora_weights.safetensors"
        echo "  Checkpoints: $SCRIPT_DIR/outputs/checkpoint-*/"
        echo ""
        echo "Next step: Generate deck visuals (option 3)"
        ;;
    
    3)
        if [ ! -f "$SCRIPT_DIR/outputs/lora_weights.safetensors" ]; then
            echo "ERROR: No trained LoRA found at outputs/lora_weights.safetensors"
            echo "Please train the model first (option 2)."
            exit 1
        fi
        
        echo ""
        echo "Generating deck visuals..."
        python generate_deck_visual.py --batch --lora outputs/lora_weights.safetensors
        
        echo ""
        echo "✓ Deck visuals complete!"
        echo "  Output: $SCRIPT_DIR/assets/decks/visuals/"
        ;;
    
    4)
        echo ""
        echo "Running full pipeline..."
        echo ""
        echo "Step 1/3: Generating dataset..."
        python generate_dataset.py --config config.json
        
        echo ""
        echo "Step 2/3: Training LoRA..."
        python train_flux_lora.py --config config.json
        
        echo ""
        echo "Step 3/3: Generating deck visuals..."
        python generate_deck_visual.py --batch --lora outputs/lora_weights.safetensors
        
        echo ""
        echo "=========================================="
        echo "✓ Full pipeline complete!"
        echo "=========================================="
        echo ""
        echo "Results:"
        echo "  Dataset: $SCRIPT_DIR/dataset/"
        echo "  LoRA: $SCRIPT_DIR/outputs/lora_weights.safetensors"
        echo "  Deck visuals: $SCRIPT_DIR/assets/decks/visuals/"
        ;;
    
    5)
        read -p "Enter prompt: " prompt
        read -p "Output filename [custom.png]: " output
        output=${output:-custom.png}
        
        if [ -f "$SCRIPT_DIR/outputs/lora_weights.safetensors" ]; then
            python generate_deck_visual.py --prompt "$prompt" --output "$output" --lora outputs/lora_weights.safetensors
        else
            echo "WARNING: No LoRA found, using base model only"
            python generate_deck_visual.py --prompt "$prompt" --output "$output"
        fi
        
        echo ""
        echo "✓ Generated: $output"
        ;;
    
    *)
        echo "Invalid choice. Exiting."
        exit 1
        ;;
esac

echo ""
echo "Done!"