#!/usr/bin/env bash

set -euo pipefail

# Fixed build script with proper ROCm/HIP environment

echo "=== Fixed Build Script ==="
echo "Configuring environment for ROCm 7.0.1 with proper exports"

# Default to the ROCm 7.0.1 install unless the caller overrides it.
export ROCM_PATH="${ROCM_PATH:-/opt/rocm-7.0.1}"
export ROCM_HOME="${ROCM_HOME:-$ROCM_PATH}"
export HIP_PATH="${HIP_PATH:-$ROCM_PATH}"
export HIP_HOME="${HIP_HOME:-$ROCM_PATH}"

export PATH="$ROCM_HOME/bin:$PATH"

# Fix architecture specifications - use gfx942 consistently
export TORCH_HIP_ARCH_LIST="gfx942"
export PYTORCH_ROCM_ARCH="gfx942"

# Remove HSA_OVERRIDE_GFX_VERSION - not needed since MI300X is already gfx942
unset HSA_OVERRIDE_GFX_VERSION

# Force single-threaded compilation to avoid ninja hanging
export MAX_JOBS=1

# Enable PyTorch JIT logging for debugging
export PYTORCH_JIT_LOG_LEVEL=1

export TORCH_EXTENSIONS_DIR="${TORCH_EXTENSIONS_DIR:-$PWD/.torch_extensions}"

echo "Environment configured:"
echo "ROCM_PATH=$ROCM_PATH"
echo "TORCH_HIP_ARCH_LIST=$TORCH_HIP_ARCH_LIST"
echo "PYTORCH_ROCM_ARCH=$PYTORCH_ROCM_ARCH"
echo "MAX_JOBS=$MAX_JOBS"
echo "PYTORCH_JIT_LOG_LEVEL=$PYTORCH_JIT_LOG_LEVEL"
echo "HSA_OVERRIDE_GFX_VERSION=${HSA_OVERRIDE_GFX_VERSION:-unset}"
echo

echo "Starting build..."
python -u build.py