#!/bin/bash
# Post-build script for Python Spaces
# Builds CUDA extensions after dependencies are installed

set -e

echo "=========================================="
echo "Post-Build: Building CUDA Extensions"
echo "=========================================="

export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
export PATH=${CUDA_HOME}/bin:${PATH}
export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
export FORCE_CUDA=1
export OMP_NUM_THREADS=1

echo "CUDA_HOME: $CUDA_HOME"
echo "PyTorch version: $(python -c 'import torch; print(torch.__version__)')"
echo "CUDA available: $(python -c 'import torch; print(torch.cuda.is_available())')"

if python -c 'import torch; print(torch.cuda.is_available())' | grep -q True; then
    echo "CUDA version: $(python -c 'import torch; print(torch.version.cuda)')"
    echo "GPU: $(python -c 'import torch; print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"N/A\")')"
fi

echo "=========================================="

# Install mamba-ssm
echo ""
echo "Installing mamba-ssm..."
if python -c "import mamba_ssm" 2>/dev/null; then
    echo "mamba_ssm already installed"
else
    pip install --no-cache-dir --verbose mamba-ssm>=2.2.2 || {
        echo "Method 1 failed. Trying with build isolation disabled..."
        pip install --no-cache-dir --no-build-isolation --verbose mamba-ssm>=2.2.2 || {
            echo "WARNING: mamba_ssm installation failed. Will use fallback implementation."
            exit 1
        }
    }
    if python -c "import mamba_ssm" 2>/dev/null; then
        echo "mamba_ssm installed successfully"
    else
        echo "ERROR: mamba_ssm import failed after installation"
        exit 1
    fi
fi

# Build selective_scan_cuda_oflex
echo ""
echo "Building selective_scan_cuda_oflex..."
if python -c "import selective_scan_cuda_oflex" 2>/dev/null; then
    echo "selective_scan_cuda_oflex already installed"
else
    if [ -d "SRMA-Mamba/selective_scan" ]; then
        cd SRMA-Mamba/selective_scan
        pip install -e . || {
            echo "Method 1 (pip install -e) failed. Trying setup.py build_ext..."
            python setup.py build_ext --inplace || {
                echo "ERROR: selective_scan_cuda_oflex build failed. Will use fallback implementation."
                cd ../..
                exit 1
            }
        }
        if python -c "import selective_scan_cuda_oflex" 2>/dev/null; then
            echo "selective_scan_cuda_oflex built successfully"
            cd ../..
        else
            echo "ERROR: selective_scan_cuda_oflex import failed after build"
            cd ../..
            exit 1
        fi
    else
        echo "ERROR: SRMA-Mamba/selective_scan directory not found"
        exit 1
    fi
fi

echo ""
echo "=========================================="
echo "Post-Build: Verification"
echo "=========================================="
python -c "
import sys
try:
    import mamba_ssm
    try:
        version = mamba_ssm.__version__
        print(f'mamba_ssm: INSTALLED (version: {version})')
    except:
        print('mamba_ssm: INSTALLED')
except ImportError as e:
    print(f'mamba_ssm: NOT INSTALLED - {e}')
    sys.exit(1)

try:
    import selective_scan_cuda_oflex
    print('selective_scan_cuda_oflex: INSTALLED')
except ImportError as e:
    print(f'selective_scan_cuda_oflex: NOT INSTALLED - {e}')
    sys.exit(1)

print('All CUDA extensions verified successfully.')
"

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