import React, { useEffect, useRef } from 'react'
import type { CellParams } from '../types'

export type GridCellFocusHealProps = {
    hasFocus: boolean
    CellRenderer: React.ComponentType<CellParams<any>>
    cellRendererProps: CellParams<any>
    parentRef: React.MutableRefObject<HTMLElement | null>
}

export const GridCellFocusHeal = ({
    CellRenderer,
    cellRendererProps,
    parentRef,
    hasFocus,
}: GridCellFocusHealProps) => {
    const hadFocusBeforeRender = useRef(false)

    if (parentRef.current) {
        hadFocusBeforeRender.current = parentRef.current.contains(
            document.activeElement
        )
    }

    useEffect(() => {
        if (parentRef.current && hadFocusBeforeRender.current) {
            /*
                Chrome will consistently fall back to document.body whereas Firefox/Safari will fallback
                to another element in the same parent if possible
            */
            const needsFocusHeal =
                hasFocus &&
                (document.activeElement === document.body ||
                    (parentRef.current.contains(document.activeElement) &&
                        document.activeElement?.getAttribute('tabindex') !==
                            '0'))

            if (needsFocusHeal) {
                parentRef.current
                    .querySelector<HTMLElement>(`[tabindex="0"]`)
                    ?.focus({ preventScroll: true })
            }
        }
    })

    return <CellRenderer {...cellRendererProps} />
}