#!/bin/bash

# Exit on error
set -e

# Print info
echo "Building C++ Tracker"

# Set CUDA environment if needed
CUDA_AVAILABLE=0
if [ -z "$CUDA_HOME" ]; then
    if [ -d "/usr/local/cuda" ]; then
        export CUDA_HOME=/usr/local/cuda
        CUDA_AVAILABLE=1
    elif [ -d "/usr/lib/cuda" ]; then
        export CUDA_HOME=/usr/lib/cuda
        CUDA_AVAILABLE=1
    fi
fi

# Add CUDA to path if available
if [ $CUDA_AVAILABLE -eq 1 ]; then
    echo "CUDA found at $CUDA_HOME"
    export PATH=$CUDA_HOME/bin:$PATH
    export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

    # Determine CUDA version
    CUDA_VERSION=$(nvcc --version | grep "release" | awk '{print $6}' | cut -c2- | cut -d'.' -f1-2)
    echo "Detected CUDA version: $CUDA_VERSION"
else
    echo "CUDA not found, building in CPU-only mode"
fi

# Download and extract LibTorch with appropriate CUDA support if not already present
LIBTORCH_DIR="$HOME/libtorch"
if [ ! -d "$LIBTORCH_DIR" ]; then
    echo "Downloading LibTorch..."
    # Use a compatible version based on detected CUDA
    if [ $CUDA_AVAILABLE -eq 1 ]; then
        echo "Downloading CUDA-enabled LibTorch"
        if [[ "$CUDA_VERSION" == "11.5" || "$CUDA_VERSION" == "11.6" || "$CUDA_VERSION" == "11.7" ]]; then
            LIBTORCH_URL="https://download.pytorch.org/libtorch/cu116/libtorch-cxx11-abi-shared-with-deps-1.13.0%2Bcu116.zip"
        elif [[ "$CUDA_VERSION" == "11.3" || "$CUDA_VERSION" == "11.4" ]]; then
            LIBTORCH_URL="https://download.pytorch.org/libtorch/cu113/libtorch-cxx11-abi-shared-with-deps-1.12.1%2Bcu113.zip"
        else
            LIBTORCH_URL="https://download.pytorch.org/libtorch/cu118/libtorch-cxx11-abi-shared-with-deps-2.0.0%2Bcu118.zip"
        fi
    else
        echo "Downloading CPU-only LibTorch"
        LIBTORCH_URL="https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.0.0%2Bcpu.zip"
    fi
    
    wget $LIBTORCH_URL -O libtorch.zip
    echo "Extracting LibTorch..."
    mkdir -p $HOME
    unzip -q libtorch.zip -d $HOME
    rm libtorch.zip
    echo "LibTorch extracted to $LIBTORCH_DIR"
else
    echo "Using existing LibTorch installation at $LIBTORCH_DIR"
fi

# Create build directory
mkdir -p build
cd build

# Create local bin directory
mkdir -p ../bin

# Configure with CMake
echo "Configuring with CMake..."
if [ $CUDA_AVAILABLE -eq 1 ]; then
    cmake .. -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=.. -DCPU_ONLY=OFF
else
    cmake .. -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=.. -DCPU_ONLY=ON
fi

# Build the project
echo "Building the project..."
cmake --build . --config Release -j $(nproc)

# Install to local directory
echo "Installing to local bin directory..."
cmake --install . --config Release

echo "Build complete! Executable is in bin/"