#!/bin/bash

# Target LibTorch 1.8.0+cu111 to match Python environment
LIBTORCH_VERSION="1.8.0"
LIBTORCH_CUDA_VERSION_SHORT="cu111" # e.g., cu111, cu117, cu118
TARGET_CUDA_VERSION_FOR_LIBTORCH="11.1" # e.g., 11.1, 11.7, 11.8

# --- CUDA Setup ---
# Try to find the exact CUDA version LibTorch was built against
PROPOSED_CUDA_HOME="/usr/local/cuda-${TARGET_CUDA_VERSION_FOR_LIBTORCH}"

if [ -d "${PROPOSED_CUDA_HOME}" ] && [ -x "${PROPOSED_CUDA_HOME}/bin/nvcc" ]; then
    echo "Found targeted CUDA version for LibTorch: ${TARGET_CUDA_VERSION_FOR_LIBTORCH} at ${PROPOSED_CUDA_HOME}"
    export CUDA_HOME="${PROPOSED_CUDA_HOME}"
else
    echo "Warning: Targeted CUDA version ${TARGET_CUDA_VERSION_FOR_LIBTORCH} for LibTorch not found at ${PROPOSED_CUDA_HOME}."
    # Fallback to user's specified CUDA_HOME or a default, and warn about potential mismatch
    if [ -z "$CUDA_HOME" ]; then # If CUDA_HOME is not already set in the environment
        export CUDA_HOME=/usr/local/cuda-11.8 # Default fallback
        echo "Warning: Using fallback CUDA_HOME: $CUDA_HOME. This might mismatch LibTorch's CUDA version ${TARGET_CUDA_VERSION_FOR_LIBTORCH}."
    else
        echo "Warning: Using externally set CUDA_HOME: $CUDA_HOME. Ensure this is compatible with LibTorch for CUDA ${TARGET_CUDA_VERSION_FOR_LIBTORCH}."
    fi
    # We proceed, but there's a higher chance of issues if nvcc version doesn't align with LibTorch's CUDA build.
fi

export PATH=${CUDA_HOME}/bin:${PATH}
export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
export CMAKE_CUDA_COMPILER=${CUDA_HOME}/bin/nvcc

# Exit on error
set -e

# --- vcpkg setup ---
# Remove explicit VCPKG_INSTALLATION_ROOT.
# The CMAKE_TOOLCHAIN_FILE will be set in the cmake command if vcpkg is part of the project.
# It's often good practice to have vcpkg as a submodule or managed by CMake FetchContent.

echo "Building C++ Tracker"
echo "Using CUDA_HOME: $CUDA_HOME"
echo "Using PATH: $PATH"
echo "Using LD_LIBRARY_PATH: $LD_LIBRARY_PATH"

# Verify CUDA version from specified CUDA_HOME
NVCC_PATH="${CUDA_HOME}/bin/nvcc"
if ! [ -x "${NVCC_PATH}" ]; then
  echo "Error: NVCC not found at specified path: ${NVCC_PATH}" >&2
  echo "Please ensure CUDA_HOME is set correctly and NVCC is executable." >&2
  exit 1
fi
DETECTED_CUDA_RUNTIME_VERSION=$(${NVCC_PATH} --version | grep "release" | awk '{print $6}' | cut -c2- | cut -d'.' -f1-2)
echo "Detected CUDA Runtime version (from ${NVCC_PATH}): $DETECTED_CUDA_RUNTIME_VERSION"

# Target LibTorch 1.8.0+cu111 to match Python environment
LIBTORCH_VARIANT="cxx11-abi-shared-with-deps" # For shared libraries with dependencies
LIBTORCH_DIR_NAME="libtorch_${LIBTORCH_VERSION}_${LIBTORCH_CUDA_VERSION_SHORT}"
LIBTORCH_INSTALL_PATH="$HOME/${LIBTORCH_DIR_NAME}/libtorch"
LIBTORCH_ZIP_URL="https://download.pytorch.org/libtorch/${LIBTORCH_CUDA_VERSION_SHORT}/libtorch-${LIBTORCH_VARIANT}-${LIBTORCH_VERSION}%2B${LIBTORCH_CUDA_VERSION_SHORT}.zip"
EXPECTED_LIBTORCH_CMAKE_CONFIG="${LIBTORCH_INSTALL_PATH}/share/cmake/Torch/TorchConfig.cmake"

echo "Targeting LibTorch ${LIBTORCH_VERSION}+${LIBTORCH_CUDA_VERSION_SHORT}"

if [ ! -f "$EXPECTED_LIBTORCH_CMAKE_CONFIG" ]; then
    echo "LibTorch ${LIBTORCH_VERSION}+${LIBTORCH_CUDA_VERSION_SHORT} not found at ${LIBTORCH_INSTALL_PATH}."
    TMP_DIR=$(mktemp -d)
    echo "Downloading LibTorch ${LIBTORCH_VERSION}+${LIBTORCH_CUDA_VERSION_SHORT}..."
    echo "Using LibTorch URL: ${LIBTORCH_ZIP_URL}"
    wget -q -O "${TMP_DIR}/libtorch.zip" "${LIBTORCH_ZIP_URL}"
    echo "Extracting LibTorch to $HOME/${LIBTORCH_DIR_NAME}..."
    mkdir -p "$HOME/${LIBTORCH_DIR_NAME}"
    unzip -q "${TMP_DIR}/libtorch.zip" -d "$HOME/${LIBTORCH_DIR_NAME}"
    rm -rf "${TMP_DIR}"
    if [ -f "$EXPECTED_LIBTORCH_CMAKE_CONFIG" ]; then
        echo "LibTorch ${LIBTORCH_VERSION}+${LIBTORCH_CUDA_VERSION_SHORT} extracted to ${LIBTORCH_INSTALL_PATH}"
    else
        echo "Error: LibTorch extraction failed or TorchConfig.cmake not found at expected location: $EXPECTED_LIBTORCH_CMAKE_CONFIG"
        exit 1
    fi
else
    echo "Found existing LibTorch ${LIBTORCH_VERSION}+${LIBTORCH_CUDA_VERSION_SHORT} at ${LIBTORCH_INSTALL_PATH}"
fi

# Set Torch_DIR for CMake
export Torch_DIR="${LIBTORCH_INSTALL_PATH}/share/cmake/Torch"

# Build directory
BUILD_DIR="build"
mkdir -p "${BUILD_DIR}"
cd "${BUILD_DIR}"

# Relative path to vcpkg.cmake from the build directory if vcpkg is a submodule in the project root
# This assumes vcpkg is at ../vcpkg (relative to build dir) or ./vcpkg (relative to project root)
VCPKG_TOOLCHAIN_FILE_PROJECT_SUBMODULE="../vcpkg/scripts/buildsystems/vcpkg.cmake"
VCPKG_TOOLCHAIN_FILE_BUILD_SUBDIR="vcpkg/scripts/buildsystems/vcpkg.cmake" # if vcpkg is cloned into build dir by CMake

CMAKE_TOOLCHAIN_ARG=""
if [ -f "$VCPKG_TOOLCHAIN_FILE_PROJECT_SUBMODULE" ]; then
    echo "Using vcpkg toolchain: $VCPKG_TOOLCHAIN_FILE_PROJECT_SUBMODULE (relative to build dir)"
    CMAKE_TOOLCHAIN_ARG="-DCMAKE_TOOLCHAIN_FILE=$VCPKG_TOOLCHAIN_FILE_PROJECT_SUBMODULE"
elif [ -f "$VCPKG_TOOLCHAIN_FILE_BUILD_SUBDIR" ]; then # this case is less common for user setup
    echo "Using vcpkg toolchain: $VCPKG_TOOLCHAIN_FILE_BUILD_SUBDIR (relative to build dir)"
    CMAKE_TOOLCHAIN_ARG="-DCMAKE_TOOLCHAIN_FILE=$VCPKG_TOOLCHAIN_FILE_BUILD_SUBDIR"
else
    echo "vcpkg.cmake not found at common project submodule paths. CMake will try to find it via vcpkg.json if vcpkg is globally installed and configured."
    # If vcpkg is globally installed and CMAKE_TOOLCHAIN_FILE is set in user's environment, CMake might pick it up.
    # Or, if CMake has native vcpkg integration via CMAKE_PROJECT_TOP_LEVEL_INCLUDES with vcpkg.cmake.
fi

echo "Configuring with CMake..."
# Pass the CMAKE_TOOLCHAIN_FILE to cmake.
# The Torch_DIR is already exported.
cmake .. -DCMAKE_BUILD_TYPE=Release ${CMAKE_TOOLCHAIN_ARG}

echo "Building..."
cmake --build . --config Release -- -j$(nproc)

echo "Build complete."
cd ..

echo "Build complete! Executable is in bin/"