#!/usr/bin/env bash
# ==============================================================================
# CUDASlicer Automated Build & Installation Script
# ==============================================================================
# Supports Linux and WSL environments with CUDA GPU acceleration.
#
# Usage:
#   ./build.sh [OPTIONS]
#
# Options:
#   --clean         Clean all build artifacts before building
#   --test          Run test suite and benchmarks after building
#   --install       Install Python package into active Python environment
#   --c-only        Build only standalone C/C++ libraries and binaries
#   --debug         Build with Debug symbols (-g / -O0)
#   --release       Build with Release optimizations (-O3, default)
#   -j, --jobs N    Number of parallel compilation jobs (default: all cores)
#   -h, --help      Show this help message
# ==============================================================================

set -e

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

# Color formatting
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# Default parameters
DO_CLEAN=false
DO_TEST=false
DO_INSTALL=false
DO_WHEEL=true
BUILD_TYPE="Release"
BUILD_C_ONLY=false
NUM_JOBS=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case "$1" in
        --clean)
            DO_CLEAN=true
            shift
            ;;
        --test)
            DO_TEST=true
            shift
            ;;
        --install)
            DO_INSTALL=true
            shift
            ;;
        --wheel)
            DO_WHEEL=true
            shift
            ;;
        --no-wheel)
            DO_WHEEL=false
            shift
            ;;
        --c-only)
            BUILD_C_ONLY=true
            DO_WHEEL=false
            shift
            ;;
        --debug)
            BUILD_TYPE="Debug"
            shift
            ;;
        --release)
            BUILD_TYPE="Release"
            shift
            ;;
        -j|--jobs)
            NUM_JOBS="$2"
            shift 2
            ;;
        -h|--help)
            echo "CUDASlicer Build Script"
            echo "Usage: ./build.sh [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --clean         Remove build directory, libraries, and binaries first"
            echo "  --test          Run unit tests and benchmarks upon completion"
            echo "  --install       Install Python package into active environment via pip"
            echo "  --wheel         Build standalone Python .whl distribution package in dist/ (default: enabled)"
            echo "  --no-wheel      Skip building Python .whl distribution package"
            echo "  --c-only        Build only C/C++ libraries (lib/ and bin/)"
            echo "  --debug         Compile with Debug flags"
            echo "  --release       Compile with Release flags (default)"
            echo "  -j, --jobs N    Parallel build jobs (default: ${NUM_JOBS})"
            echo "  -h, --help      Show this help message"
            exit 0
            ;;
        *)
            echo -e "${RED}[Error] Unknown option: $1${NC}"
            echo "Run './build.sh --help' for usage."
            exit 1
            ;;
    esac
done

echo -e "${BOLD}${CYAN}========================================================================${NC}"
echo -e "${BOLD}${CYAN} CUDASlicer: High-Performance GPU Slicing & CSG Engine Build${NC}"
echo -e "${BOLD}${CYAN}========================================================================${NC}"

# Check for required tools
echo -e "\n${BLUE}[1/5] Checking Build Environment & Dependencies...${NC}"

if command -v nvcc &>/dev/null; then
    NVCC_VERSION=$(nvcc --version | grep "release" | awk '{print $5}' | tr -d ',')
    echo -e "  * NVIDIA CUDA Compiler (NVCC): ${GREEN}Found (v${NVCC_VERSION})${NC}"
else
    echo -e "  * NVIDIA CUDA Compiler (NVCC): ${YELLOW}Not Found (CPU fallback only)${NC}"
fi

if command -v cmake &>/dev/null; then
    CMAKE_VERSION=$(cmake --version | head -n 1 | awk '{print $3}')
    echo -e "  * CMake:                      ${GREEN}Found (v${CMAKE_VERSION})${NC}"
else
    echo -e "  * CMake:                      ${RED}Missing! Please install cmake (sudo apt install cmake)${NC}"
    exit 1
fi

if command -v python3 &>/dev/null; then
    PY_VERSION=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}')")
    echo -e "  * Python:                     ${GREEN}Found (v${PY_VERSION})${NC}"
else
    echo -e "  * Python:                     ${YELLOW}python3 not found in PATH${NC}"
fi

# Detect GPU Architecture if available
CUDA_ARCHS="75;80;86"
if command -v nvidia-smi &>/dev/null; then
    GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -n 1 || true)
    if [[ -n "${GPU_NAME}" ]]; then
        echo -e "  * Detected GPU Device:        ${GREEN}${GPU_NAME}${NC}"
    fi
fi

# Handle clean
if [ "$DO_CLEAN" = true ]; then
    echo -e "\n${BLUE}[2/5] Cleaning previous build artifacts...${NC}"
    rm -rf build/ bin/ lib/ dist/ python/cudaslicer/*.so python/cudaslicer.egg-info
    echo -e "  * ${GREEN}Clean complete.${NC}"
else
    echo -e "\n${BLUE}[2/5] Skipping clean (incremental build). Pass --clean to reset.${NC}"
fi

# Configure with CMake
echo -e "\n${BLUE}[3/5] Configuring CMake (${BUILD_TYPE})...${NC}"
mkdir -p build

CMAKE_EXTRA_ARGS=()
if [ "$BUILD_C_ONLY" = true ]; then
    CMAKE_EXTRA_ARGS+=("-DBUILD_PYTHON_BINDINGS=OFF")
fi

cmake -B build -S . \
    -DCMAKE_BUILD_TYPE="${BUILD_TYPE}" \
    -DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHS}" \
    "${CMAKE_EXTRA_ARGS[@]}"

# Build target binaries
echo -e "\n${BLUE}[4/5] Compiling targets with ${NUM_JOBS} parallel workers...${NC}"
cmake --build build --config "${BUILD_TYPE}" -j "${NUM_JOBS}"

echo -e "\n${GREEN}[SUCCESS] Compilation finished successfully!${NC}"
echo -e "  * Shared Library:  ${CYAN}lib/libcudaslicer_offset.so${NC}"
echo -e "  * Static Library:  ${CYAN}lib/libcudaslicer_offset.a${NC}"
echo -e "  * Benchmark Exec:  ${CYAN}bin/gpu_slicer_c_bench${NC}"
if [ "$BUILD_C_ONLY" = false ]; then
    echo -e "  * Python Backend:  ${CYAN}python/cudaslicer/_cudaslicer_backend*.so${NC}"
fi

# Python Wheel Package Build
if [ "$DO_WHEEL" = true ]; then
    echo -e "\n${BLUE}[5/5] Building Python Wheel package (.whl)...${NC}"
    mkdir -p dist
    python3 setup.py bdist_wheel
    WHL_FILE=$(ls -t dist/*.whl 2>/dev/null | head -n 1 || true)
    if [ -n "${WHL_FILE}" ]; then
        echo -e "  * ${GREEN}Built Wheel Package: ${CYAN}${WHL_FILE}${NC}"
    fi
fi

# Optional Python Installation
if [ "$DO_INSTALL" = true ]; then
    echo -e "\n${BLUE}[Optional] Installing Python package in editable mode...${NC}"
    pip install -e . --no-build-isolation
    echo -e "  * ${GREEN}Package 'cudaslicer' installed into current environment.${NC}"
fi

# Optional Tests & Benchmarks
if [ "$DO_TEST" = true ]; then
    echo -e "\n${BLUE}[Test Suite] Running Tests & Benchmarks...${NC}"
    
    echo -e "\n--- [A] Running Python GPU SDF & CSG Unit Tests ---"
    python3 tests_and_tools/test_gpu_sdf.py

    echo -e "\n--- [B] Running GPU SDF CSG Example Pipeline ---"
    python3 examples/gpu_sdf_csg_example.py

    echo -e "\n--- [C] Running C++ Slicing Benchmark ---"
    if [ -f "benchy.stl" ]; then
        ./bin/gpu_slicer_c_bench benchy.stl
    fi
fi

echo -e "\n${BOLD}${GREEN}========================================================================${NC}"
echo -e "${BOLD}${GREEN} Build completed successfully!${NC}"
echo -e "${BOLD}${GREEN}========================================================================${NC}\n"