#!/usr/bin/env bash
set -e

echo "=========================================="
echo "Post-Build: Installing CUDA Extensions"
echo "=========================================="
echo "Note: This may fail on Python Spaces without nvcc."
echo "For guaranteed success, switch to Docker Space."
echo "=========================================="

if ! command -v nvcc &> /dev/null; then
    echo "ERROR: nvcc (CUDA compiler) not found!"
    echo "Python Spaces do not provide CUDA toolchain."
    echo "ACTION REQUIRED: Switch to Docker Space in HF Settings -> Runtime -> Docker"
    echo "Docker Space includes nvcc and can build CUDA extensions."
    exit 1
fi

echo "nvcc found: $(nvcc --version | head -n 1)"
pip install pybind11 || {
    echo "WARNING: pybind11 installation failed"
}

echo ""
echo "Installing mamba-ssm..."
pip install "mamba-ssm>=2.2.2" || {
    echo "ERROR: mamba-ssm installation failed. This requires CUDA toolchain (nvcc)."
    echo "Switch to Docker Space for guaranteed CUDA extension builds."
    exit 1
}

if [ -d "SRMA-Mamba/selective_scan" ]; then
    echo ""
    echo "Building selective_scan_cuda_oflex..."
    cd SRMA-Mamba/selective_scan
    pip install -e . || {
        echo "ERROR: selective_scan_cuda_oflex build failed. This requires CUDA toolchain (nvcc)."
        echo "Switch to Docker Space for guaranteed CUDA extension builds."
        cd ../..
        exit 1
    }
    cd ../..
else
    echo "ERROR: SRMA-Mamba/selective_scan directory not found"
    exit 1
fi

echo ""
echo "Verifying CUDA extensions..."
python - <<'PY'
import importlib
import mamba_ssm
importlib.import_module("selective_scan_cuda_oflex")
print("CUDA extensions OK")
PY

if [ $? -ne 0 ]; then
    echo "ERROR: CUDA extension import test failed."
    echo "Switch to Docker Space for guaranteed CUDA extension builds."
    exit 1
fi

echo "=========================================="
echo "Post-Build: Complete"
echo "=========================================="