#!/usr/bin/env bash
set -euo pipefail

RELEASE_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RUNTIME_ROOT="${RELEASE_ROOT}/runtime"
PYTHON_BIN="${PYTHON_BIN:-python}"
CUDA_TOOLKIT_ROOT="${CUDA_HOME:-/usr/local/cuda}"
NVCC="${CUDA_TOOLKIT_ROOT}/bin/nvcc"

if [[ ! -x "${NVCC}" ]]; then
  echo "nvcc was not found at ${NVCC}" >&2
  exit 1
fi

mkdir -p "${RUNTIME_ROOT}"

"${NVCC}" \
  -std=c++17 -O3 -arch=sm_120a -shared -Xcompiler=-fPIC \
  "${RUNTIME_ROOT}/nvfp4_linear.cu" \
  -o "${RUNTIME_ROOT}/libmage_nvfp4_linear.so" \
  -L"${CUDA_TOOLKIT_ROOT}/lib64" -lcublasLt -lcublas

PYTORCH_INCLUDE_FLAGS="$("${PYTHON_BIN}" - <<'PY'
from torch.utils.cpp_extension import include_paths
print(" ".join(f"-I{path}" for path in include_paths()))
PY
)"
PYTORCH_LIBRARY_FLAGS="$("${PYTHON_BIN}" - <<'PY'
from torch.utils.cpp_extension import library_paths
print(" ".join(f"-L{path}" for path in library_paths()))
PY
)"
PYTORCH_RPATH_FLAGS="$("${PYTHON_BIN}" - <<'PY'
from torch.utils.cpp_extension import library_paths
print(" ".join(f"-Wl,-rpath,{path}" for path in library_paths()))
PY
)"
PYTORCH_ABI_FLAG="$("${PYTHON_BIN}" - <<'PY'
import torch
print(f"-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}")
PY
)"

g++ \
  -std=c++20 -O3 -shared -fPIC \
  ${PYTORCH_ABI_FLAG} \
  ${PYTORCH_INCLUDE_FLAGS} \
  ${PYTORCH_LIBRARY_FLAGS} \
  ${PYTORCH_RPATH_FLAGS} \
  -I"${CUDA_TOOLKIT_ROOT}/include" \
  -I"${RUNTIME_ROOT}" \
  -o "${RUNTIME_ROOT}/libmage_nvfp4_torch_op.so" \
  "${RUNTIME_ROOT}/sm120_linear_op.cpp" \
  -L"${CUDA_TOOLKIT_ROOT}/lib64" \
  -L"${RUNTIME_ROOT}" \
  -Wl,-rpath,"${CUDA_TOOLKIT_ROOT}/lib64" \
  -Wl,-rpath,'$ORIGIN' \
  -ltorch -ltorch_cpu -ltorch_cuda -lc10 -lc10_cuda -lcudart \
  -lmage_nvfp4_linear

echo "Built the packaged SM120 runtime in ${RUNTIME_ROOT}"