#!/bin/bash
# RustyTorch++ Development Script with GPU Acceleration

set -e

# Source environment for CUDA
source ~/.zshrc 2>/dev/null || source ~/.bashrc 2>/dev/null

# Configuration
RUSTG_PATH="../rust/rustg"
CARGO_G="$RUSTG_PATH/cargo-g/target/release/cargo-g"
GPU_DEV_TOOLS="$RUSTG_PATH/gpu-dev-tools/target/release/gpu-dev-tools"

echo "=== RustyTorch++ GPU-Accelerated Development ==="
echo "GPU: $(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null || echo 'Not detected')"
echo "CUDA: $(nvcc --version 2>/dev/null | grep release | awk '{print $5}' | sed 's/,//' || echo 'Not available')"

# Function to use cargo-g if available, fallback to cargo
cargo_gpu() {
    if [ -x "$CARGO_G" ]; then
        echo "🚀 Using cargo-g for GPU-accelerated build..."
        "$CARGO_G" "$@"
    elif [ -x "$RUSTG_PATH/mock-cargo-g.sh" ]; then
        echo "🔧 Using mock cargo-g..."
        "$RUSTG_PATH/mock-cargo-g.sh" "$@"
    else
        echo "📦 Using standard cargo..."
        cargo "$@"
    fi
}

# Function to format code with GPU acceleration
format_gpu() {
    if [ -x "$GPU_DEV_TOOLS" ]; then
        echo "⚡ GPU-accelerated formatting..."
        "$GPU_DEV_TOOLS" format --write crates/
    else
        echo "📝 Standard formatting..."
        cargo fmt
    fi
}

# Function to lint with GPU acceleration
lint_gpu() {
    if [ -x "$GPU_DEV_TOOLS" ]; then
        echo "⚡ GPU-accelerated linting..."
        "$GPU_DEV_TOOLS" lint crates/
    else
        echo "🔍 Standard linting..."
        cargo clippy -- -D warnings
    fi
}

# Parse command
case "${1:-build}" in
    build)
        cargo_gpu build --release
        ;;
    test)
        cargo_gpu test --all
        ;;
    bench)
        cargo_gpu bench
        ;;
    format)
        format_gpu
        ;;
    lint)
        lint_gpu
        ;;
    kernel)
        echo "🔥 Compiling GPU kernels..."
        ./scripts/compile_kernels.sh
        ;;
    clean)
        cargo clean
        rm -rf target/kernel_cache/
        ;;
    all)
        format_gpu
        lint_gpu
        cargo_gpu build --release
        cargo_gpu test --all
        ;;
    *)
        echo "Usage: $0 {build|test|bench|format|lint|kernel|clean|all}"
        exit 1
        ;;
esac

echo "=== Development task complete ==="