StaggeredGridLayoutManager with auto-fit span count

I need a StaggeredGridLayoutManager that follows these two requirements:

• Automatically calculate the number of columns based on a set minimum dp size for each column, with columns expanding to perfectly fit the available space (e.g: if I define 150dp minimum column with, and the app runs on a 480dp width device, the grid view will contain 3 columns, and the remaining 30dp will be split evenly, so that each column is 160dp wide)
• Grid layout has to recalculate span count on orientation change, as the number of columns that fit the screen in portrait and landscape mode are usually quite far apart

Until now, I haven't had the need for the grid items to be staggered, so I used an extension of GridLayoutManager to provide this functionality:

import androidx.recyclerview.widget.GridLayoutManager
import androidx.recyclerview.widget.RecyclerView
import androidx.recyclerview.widget.RecyclerView.Recycler
import kotlin.math.max
import kotlin.math.min

class GridAutoFitLayoutManager : GridLayoutManager {
    private var mColumnWidth = 0
    private var mMaximumColumns: Int
    private var mLastCalculatedWidth = -1

    @JvmOverloads
    constructor(
        context: Context?,
        columnWidthDp: Int,
        maxColumns: Int = 99
    ) : super(context, 1) //Initially set spanCount to 1, will be changed automatically later.
    {
        mMaximumColumns = maxColumns
        setColumnWidth(columnWidthDp)
    }

    private fun setColumnWidth(newColumnWidth: Int) {
        if (newColumnWidth > 0 && newColumnWidth != mColumnWidth) {
            mColumnWidth = newColumnWidth
        }
    }

    override fun onLayoutChildren(
        recycler: Recycler,
        state: RecyclerView.State
    ) {
        val width = width
        val height = height
        if (width != mLastCalculatedWidth && mColumnWidth > 0 && width > 0 && height > 0) {
            val totalSpace: Int = if (orientation == RecyclerView.VERTICAL) {
                width - paddingRight - paddingLeft
            } else {
                height - paddingTop - paddingBottom
            }
            val spanCount = min(
                mMaximumColumns,
                max(1, totalSpace / mColumnWidth)
            )
            setSpanCount(spanCount)
            mLastCalculatedWidth = width
        }
        super.onLayoutChildren(recycler, state)
    }
}

Using the layout manager in the recycler view:

recyclerView.layoutManager = GridAutoFitLayoutManager(context, resources.getDimension(R.dimen.grid_column_width).toInt())

Lastly, the item view layouts are set to match_parent on width, so they'll occupy as much space as the layout manager will let them.

However, I'm having trouble getting this to work with StaggeredGridLayoutManager. Here's my code:

import android.content.Context
import android.widget.LinearLayout
import androidx.recyclerview.widget.RecyclerView
import androidx.recyclerview.widget.RecyclerView.Recycler
import androidx.recyclerview.widget.StaggeredGridLayoutManager
import kotlin.math.max
import kotlin.math.min

class GridAutoFitStaggeredLayoutManager : StaggeredGridLayoutManager {
    private var mColumnWidth = 0
    private var mMaximumColumns: Int
    private var mLastCalculatedWidth = -1

    @JvmOverloads
    constructor(
        context: Context?,
        columnWidthDp: Int,
        maxColumns: Int = 99
    ) : super(1, LinearLayout.VERTICAL) //Initially set spanCount to 1, will be changed automatically later.
    {
        mMaximumColumns = maxColumns
        setColumnWidth(columnWidthDp)
    }

    private fun setColumnWidth(newColumnWidth: Int) {
        if (newColumnWidth > 0 && newColumnWidth != mColumnWidth) {
            mColumnWidth = newColumnWidth
        }
    }

    override fun onLayoutChildren(
        recycler: Recycler,
        state: RecyclerView.State
    ) {
        val width = width
        val height = height
        if (width != mLastCalculatedWidth && mColumnWidth > 0 && width > 0 && height > 0) {
            val totalSpace: Int = if (orientation == RecyclerView.VERTICAL) {
                width - paddingRight - paddingLeft
            } else {
                height - paddingTop - paddingBottom
            }
            val spanCount = min(
                mMaximumColumns,
                max(1, totalSpace / mColumnWidth)
            )
            setSpanCount(spanCount)
            mLastCalculatedWidth = width
        }
        super.onLayoutChildren(recycler, state)
    }
}

As soon as the recycler view is loaded, I get the following exception:

java.lang.IllegalStateException: Cannot call this method while RecyclerView is computing a layout or scrolling

The issue is because I'm calling setSpanCount() while the layout is resolving, which wasn't a problem with GridLayoutManager, but not allowed in StaggeredGridLayoutManager. My question is, if I can't set the span count during onLayoutChildren(), when can I set it?

I'm aware I can just calculate how many columns will fit prior to initializing the layout manager, and pass the span count to the StaggeredGridLayoutManager constructor. However, that won't react to orientation changes, unless I add additional logic to the activity that re-creates the layout manager, which I'd prefer to avoid. Is there a way to make my GridAutoFitStaggeredLayoutManager calculate span count and handle orientation changes on its own?