#!/bin/bash
# THE BAKER — Environment Setup Script
# 
# This script sets up the Python environment for Flux.2 LoRA training.

set -e

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

echo "=========================================="
echo "THE BAKER — Environment Setup"
echo "=========================================="
echo ""

# Find Python
PYTHON=""

# Try system Python 3.13
if [ -f "/nix/store/ppwq97q6p2vqrma5rrw74l9v1ggmlmgx-python3-3.13.12-env/bin/python3.13" ]; then
    PYTHON="/nix/store/ppwq97q6p2vqrma5rrw74l9v1ggmlmgx-python3-3.13.12-env/bin/python3.13"
    echo "✓ Found Python 3.13 in nix store"
# Try system Python 3.12
elif [ -f "/nix/store/2jkici78j482hva8f4fjas2zhwzlkwh3-python3-3.12.13-env/bin/python3.12" ]; then
    PYTHON="/nix/store/2jkici78j482hva8f4fjas2zhwzlkwh3-python3-3.12.13-env/bin/python3.12"
    echo "✓ Found Python 3.12 in nix store"
# Try cultguard devenv (if it exists)
elif [ -f "/home/workspaces/cultguard-agents/.devenv/state/venv/bin/python" ]; then
    PYTHON="/home/workspaces/cultguard-agents/.devenv/state/venv/bin/python"
    echo "✓ Found Python in cultguard devenv"
# Try cultshot devenv (if it exists)
elif [ -f "/home/workspaces/cultshot-radial/.devenv/state/venv/bin/python" ]; then
    PYTHON="/home/workspaces/cultshot-radial/.devenv/state/venv/bin/python"
    echo "✓ Found Python in cultshot devenv"
else
    echo "ERROR: No Python installation found."
    echo ""
    echo "Please install Python 3.12+ with pip, or create a venv:"
    echo "  python3 -m venv /home/workspaces/thebaker/training/.venv"
    echo "  source /home/workspaces/thebaker/training/.venv/bin/activate"
    exit 1
fi

echo "  Using: $PYTHON"
echo ""

# Check Python version
VERSION=$($PYTHON --version 2>&1)
echo "✓ Python version: $VERSION"
echo ""

# Check pip
echo "Checking pip..."
$PYTHON -m pip --version >/dev/null 2>&1 || {
    echo "ERROR: pip not available"
    exit 1
}
echo "✓ pip available"
echo ""

# Check GPU
echo "Checking GPU..."
if nvidia-smi >/dev/null 2>&1; then
    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader,nounits | head -1)
    GPU_MEM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | head -1)
    echo "✓ GPU: $GPU_NAME ($GPU_MEM MB)"
    
    # Check CUDA
    if $PYTHON -c "import torch; print(torch.cuda.is_available())" 2>/dev/null | grep -q "True"; then
        echo "✓ CUDA available in PyTorch"
    else
        echo "⚠ CUDA not available in PyTorch (will install)"
    fi
else
    echo "⚠ WARNING: No NVIDIA GPU detected"
    echo "  Training will fail without GPU acceleration"
fi
echo ""

# Check existing packages
echo "Checking installed packages..."
TORCH_INSTALLED=$($PYTHON -c "import torch; print(torch.__version__)" 2>/dev/null || echo "")
DIFFUSERS_INSTALLED=$($PYTHON -c "import diffusers; print(diffusers.__version__)" 2>/dev/null || echo "")

if [ -n "$TORCH_INSTALLED" ] && [ -n "$DIFFUSERS_INSTALLED" ]; then
    echo "✓ torch: $TORCH_INSTALLED"
    echo "✓ diffusers: $DIFFUSERS_INSTALLED"
    echo ""
    echo "Dependencies already installed!"
    echo ""
    echo "You can now run:"
    echo "  ./00_quickstart.sh"
    exit 0
fi

echo "Some packages missing. Installing..."
echo ""

# Create local venv if needed
if [[ "$PYTHON" == *"/nix/store/"* ]]; then
    echo "Creating local virtual environment..."
    $PYTHON -m venv "$SCRIPT_DIR/.venv"
    source "$SCRIPT_DIR/.venv/bin/activate"
    PYTHON="$SCRIPT_DIR/.venv/bin/python"
    echo "✓ Local venv created at .venv/"
    echo ""
fi

# Install requirements
echo "Installing requirements..."
$PYTHON -m pip install --upgrade pip
$PYTHON -m pip install -r requirements.txt

echo ""
echo "=========================================="
echo "✓ Setup complete!"
echo "=========================================="
echo ""
echo "Environment ready for training."
echo ""
echo "Next step:"
echo "  ./00_quickstart.sh"
echo ""
echo "Or manually:"
echo "  source .venv/bin/activate  (if created)"
echo "  python generate_dataset.py --config config.json"