import React, { useEffect, useLayoutEffect } from 'react'
import { useGridComponents } from '../../context/grid-components-context/hook'
import type { Range } from '@tanstack/react-virtual'
import {
    defaultRangeExtractor,
    useVirtualizer,
    observeElementRect as observeElementRectOriginal,
} from '@tanstack/react-virtual'
import { ListWrap } from '../../layout'
import {
    useGridContext,
    useGridSelector,
} from '../../context/grid-context/hook'
import { useGridNavigation } from '../../hooks/use-grid-navigation'
import { useScrollSync } from '../../context/grid-sync-context'
import {
    size,
    useResizeObserver,
    spacing,
    usePrevious,
    useTesting,
    isTestEnvironment,
} from '@planview/pv-utilities'
import { useAriaLoadingMessage, useRangeCopy, useRangePaste } from '../../hooks'
import {
    ensureCellIsVisible,
    ensureRowIsVisible,
} from '../../utils/ensure-visible'
import { GridFooterRow, GridLiveRegion } from '../../functional'
import { useDroppable } from '@dnd-kit/core'
import { DROPPABLE_TYPE_ROW_GROUP, ROW_FOCUS_ID } from '../../constants'
import {
    getColumnOffset,
    isInViewHorizontal,
    isInViewVertical,
    scrollToRowIndex,
} from '../../utils/scroll'
import { useGridPrivateSettings } from '../../context/grid-private-settings-context'

export type GridBaseProps = {
    loading?: boolean | number
    rowHeight: number
    emptyContent?: React.ReactNode
    label?: string
    disableHorizontalScroll?: boolean
}

type ObserveElementRectCallback = (rect: {
    width: number
    height: number
}) => void

export const GridBase = ({
    label,
    rowHeight,
    loading: loadingIncoming,
    emptyContent,
    disableHorizontalScroll = false,
}: GridBaseProps) => {
    const {
        GridContainer,
        GridLayout,
        GridHeader,
        GridBody,
        GridFooter,
        GridRangeSelection,
        GridRow,
        GridDragIndicator,
        GridEmptyContainer,
        GridStickyOverlay,
        GridStickySignal,
        GridContextActionsMenu,
        GridColumnMenu,
    } = useGridComponents()

    const grid = useGridContext()
    const { events, selectors, dispatch } = grid

    useEffect(
        () =>
            events.on('onFocusUnmount', () => {
                //TODO: Handle activeElement unmounting. (Maybe row is deleted etc..)
            }),
        [events]
    )
    const [hasFocus, setHasFocus] = React.useState(false)
    const currentFocus = useGridSelector(selectors.selectCurrentFocus)
    const gridNavigationEnabled = useGridSelector(
        selectors.selectNavigationEnabled
    )
    const gridCanEdit = useGridSelector(selectors.selectGridSupportsEdit)
    const columnsCount = useGridSelector(selectors.selectColumnCount)

    const stateRowIds = useGridSelector(selectors.selectRowIds)
    const stateColumnIds = useGridSelector(selectors.selectColumnIds)
    const isTreeGrid = useGridSelector(selectors.selectIsTreeGrid)
    const isSpreadsheet = useGridSelector(selectors.selectIsSpreadsheet)
    const footerEnabled = useGridSelector(selectors.selectAggregationEnabled)
    const focusedRowIndex = React.useMemo(
        () => stateRowIds.indexOf(currentFocus.rowId),
        [stateRowIds, currentFocus.rowId]
    )
    const focusedColumnIndex = React.useMemo(
        () => stateColumnIds.indexOf(currentFocus.columnId),
        [stateColumnIds, currentFocus.columnId]
    )
    const selectionMode = useGridSelector(selectors.selectSelectionMode)
    const columnWidth = useGridSelector(selectors.selectTotalColumnWidth)
    const columnWidths = useGridSelector(selectors.selectColumnWidths)
    const headerHeight = useGridSelector(selectors.selectHeaderHeight)

    const stickyColumnDetails = useGridSelector(
        selectors.selectStickyColumnDetails
    )

    const stickyColumnIndexes = stickyColumnDetails.indexes

    const keyDownHandler = useGridNavigation()

    const parentRef = React.useRef<HTMLDivElement>(null)

    const { height } = useResizeObserver({ ref: parentRef })

    useEffect(
        () =>
            events.on('onReturnFocus', () => {
                const found =
                    parentRef.current?.querySelector<HTMLElement>(
                        '[tabindex="0"]'
                    )
                if (found) {
                    found.focus()
                }
            }),
        [events]
    )

    /* By combining incoming props with state, we can ensure loading is turned on immediately
       and only turned off when the state has been updated */
    const loading = useGridSelector(
        (state) => loadingIncoming || selectors.selectLoading(state)
    )
    const loadingIds = React.useMemo(() => {
        const count = !loading ? 0 : loading === true ? 10 : loading
        return Array.from(Array(count), (_, index) => `__loading-${index + 1}`)
    }, [loading])
    const rowIds = loading ? loadingIds : stateRowIds
    const actualRowCount = useGridSelector(selectors.selectAriaRowCount)
    const rowCount = Math.max(rowIds.length, stateRowIds.length)

    const testing = useTesting()

    /**
     * This is code to make it easier to run the virtualized grid in a JSDOM
     * environment. If rendered in a TestingProvider, the tester can override
     * the calculated dimensions of the grid.
     *
     * When in a test environment, but not wrapped in TestingProvider, then the
     * grid will be 800x600
     */
    const observeElementRect = React.useMemo(
        () =>
            testing.inTestingEnvironment
                ? (_: unknown, cb: ObserveElementRectCallback) => {
                      cb({
                          width: testing.clientWidth,
                          height: testing.clientHeight,
                      })
                  }
                : isTestEnvironment
                  ? (_: unknown, cb: ObserveElementRectCallback) => {
                        cb({
                            width: 800,
                            height: 600,
                        })
                    }
                  : observeElementRectOriginal,
        [testing]
    )

    const rowVirtualizer = useVirtualizer({
        observeElementRect,
        getScrollElement: React.useCallback(() => parentRef.current, []),
        estimateSize: React.useCallback(() => rowHeight, [rowHeight]),
        getItemKey: React.useCallback(
            (index: number) => rowIds[index],
            [rowIds]
        ),
        rangeExtractor: React.useCallback(
            (range: Range) => {
                if (!range.count) {
                    return defaultRangeExtractor(range)
                }
                const next = defaultRangeExtractor(range)
                if (focusedRowIndex === -1 || next.includes(focusedRowIndex)) {
                    return next
                }
                return [focusedRowIndex, ...next].sort((a, b) => a - b)
            },
            [focusedRowIndex]
        ),
        count: rowCount,
        overscan: 10,
    })

    const rowVirtualItems = rowVirtualizer.getVirtualItems()
    const rowTotalSize = rowVirtualizer.getTotalSize()

    React.useEffect(() => {
        const scrollTop = parentRef.current?.scrollTop || 0
        const height = parentRef.current?.getBoundingClientRect().height || 0

        if (!height) {
            return
        }

        const scrolledRows = Math.max(0, (scrollTop - size.small) / rowHeight)
        const visibleRows = Math.max(0, (height - size.small) / rowHeight)

        const scrolledFullRows = Math.floor(scrolledRows)
        const visibleFullRows = Math.ceil(
            visibleRows + scrolledRows - scrolledFullRows
        )

        events.emit('onVisibleRowsChange', [
            scrolledFullRows,
            scrolledFullRows + visibleFullRows,
        ])
    }, [events, rowVirtualItems, rowHeight])

    const columnVirtualizer = useVirtualizer({
        scrollPaddingStart: stickyColumnDetails.left.width,
        scrollPaddingEnd: stickyColumnDetails.right.width,
        observeElementRect,
        getScrollElement: React.useCallback(() => parentRef.current, []),
        estimateSize: React.useCallback(
            (index: number) => columnWidths[index],
            [columnWidths]
        ),
        rangeExtractor: React.useCallback(
            (range: Range) => {
                if (!range.count) {
                    return defaultRangeExtractor(range)
                }
                const next = defaultRangeExtractor(range)

                const updated = new Set([focusedColumnIndex, ...next])

                return (focusedColumnIndex >= 0 ? [...updated] : next)
                    .filter((id) => !stickyColumnIndexes.has(id))
                    .sort((a, b) => a - b)
            },
            [focusedColumnIndex, stickyColumnIndexes]
        ),
        count: columnWidths.length,
        overscan: 5,
        horizontal: true,
    })

    const columnVirtualItems = columnVirtualizer.getVirtualItems()
    const columnTotalSize = columnVirtualizer.getTotalSize()
    useEffect(() => {
        const numberOfRowsInView = Math.round((height - size.small) / rowHeight)
        if (numberOfRowsInView > 0) {
            dispatch({
                type: 'updateRowDimensions',
                payload: {
                    numberOfRowsInView,
                    rowHeight,
                },
            })
        }
    }, [height, rowHeight, dispatch])

    const prevColumnWidths = usePrevious(columnWidths)
    useLayoutEffect(() => {
        if (prevColumnWidths?.length !== columnWidths.length) {
            return
        }

        for (let i = 0; i < columnWidths.length; i++) {
            if (prevColumnWidths[i] !== columnWidths[i]) {
                columnVirtualizer.resizeItem(i, columnWidths[i])
            }
        }
    }, [prevColumnWidths, columnWidths, columnVirtualizer])

    useEffect(
        () =>
            events.on(
                'scrollToIndex',
                (
                    rowIndex: number | null,
                    columnIndex: number | null,
                    mode: 'auto' | 'always'
                ) => {
                    let shouldScrollRow = false

                    if (rowIndex !== null) {
                        if (mode === 'always') {
                            //Always scroll when mode is always
                            shouldScrollRow = true
                        } else {
                            const foundRow = rowVirtualItems.find(
                                (item) => item.index === rowIndex
                            )

                            shouldScrollRow =
                                !foundRow ||
                                !isInViewVertical(
                                    foundRow,
                                    parentRef.current,
                                    headerHeight,
                                    footerEnabled ? (foundRow?.size ?? 0) : 0
                                )
                        }

                        if (shouldScrollRow) {
                            /**
                             * We use our own scroll function here temporarily until the issue: https://github.com/TanStack/virtual/issues/1001
                             * is resolved in the react-virtual package.
                             */
                            // rowVirtualizer.scrollToIndex(rowIndex, {
                            //     align: 'start',
                            // })
                            scrollToRowIndex(
                                rowIndex,
                                rowHeight,
                                rowTotalSize,
                                parentRef.current,
                                headerHeight,
                                footerEnabled ? rowHeight : 0
                            )
                        }
                    }

                    if (columnIndex !== null) {
                        const columnWidths = selectors.selectColumnWidths(
                            grid.getState()
                        )

                        const leftStickyContentWidth =
                            stickyColumnDetails.left.width
                        const rightStickyContentWidth =
                            stickyColumnDetails.right.width

                        const [start, end] = getColumnOffset(
                            columnIndex,
                            columnWidths
                        )

                        if (
                            mode === 'always' ||
                            !isInViewHorizontal(
                                [start, end],
                                parentRef.current,
                                leftStickyContentWidth,
                                rightStickyContentWidth
                            )
                        ) {
                            columnVirtualizer.scrollToIndex(columnIndex)
                        }
                    }
                }
            ),
        [
            events,
            rowVirtualItems,
            rowVirtualizer,
            columnVirtualizer,
            footerEnabled,
            stickyColumnDetails,
            selectors,
            grid,
            headerHeight,
            rowHeight,
            rowTotalSize,
        ]
    )

    useScrollSync(parentRef, {
        tail: footerEnabled ? rowHeight : 0,
        head: headerHeight,
    })

    const editing = useGridSelector(selectors.selectIsEditing)

    useAriaLoadingMessage(!!loading)

    /* We want to avoid focusing the cell if the user has just clicked in the grid. */
    const pointerJustDown = React.useRef(false)

    const handlePointerDown = React.useCallback(() => {
        pointerJustDown.current = true
        setTimeout(() => {
            pointerJustDown.current = false
        }, 10)
    }, [])

    const handleKeyDown = React.useCallback(
        (e: React.KeyboardEvent) => {
            pointerJustDown.current = false
            keyDownHandler(e)
        },
        [keyDownHandler]
    )

    React.useEffect(() => {
        if (
            hasFocus &&
            !editing &&
            parentRef.current &&
            !pointerJustDown.current
        ) {
            const elementType =
                currentFocus.columnId === ROW_FOCUS_ID ? 'row' : 'cell'

            const found = parentRef.current.querySelector<HTMLElement>(
                `[data-current-${elementType}][tabindex="0"], [data-current-${elementType}] [tabindex="0"]`
            )

            if (found) {
                if (elementType === 'cell') {
                    const cell = found.closest(
                        currentFocus.area === 'header' ? 'th' : 'td'
                    )
                    if (cell) {
                        ensureCellIsVisible(cell, parentRef.current, grid)
                    }
                } else {
                    const row = found.closest('tr')
                    if (row) {
                        ensureRowIsVisible(row, parentRef.current, grid)
                    }
                }
                found.focus()
            }
        }
    }, [hasFocus, currentFocus, editing, grid])

    const { onRangeCopy, onBulkCellChange } = useGridPrivateSettings()
    const handleCopy = useRangeCopy(onRangeCopy ?? 'labels', parentRef.current)
    const handlePaste = useRangePaste(onBulkCellChange, parentRef.current)

    const renderEmpty = rowCount === 0 && !!emptyContent

    const isDraggingEnabled = useGridSelector(
        grid.selectors.selectIsDraggingEnabled
    )
    const dragNeedsOffset = useGridSelector(
        grid.selectors.selectDragNeedsFinalRowOffset
    )
    const { setNodeRef } = useDroppable({
        id: DROPPABLE_TYPE_ROW_GROUP,
        disabled: !isDraggingEnabled,
    })

    const offsetAdjustment = dragNeedsOffset ? spacing.xsmall : 0
    const bodyHeight = rowTotalSize + offsetAdjustment

    return (
        <GridContainer>
            <ListWrap
                $disableHorizontalScroll={disableHorizontalScroll}
                ref={parentRef}
                tabIndex={-1}
                onKeyDown={handleKeyDown}
                onPointerDownCapture={handlePointerDown}
                onFocus={() => {
                    setHasFocus(true)
                }}
                onBlur={(e) => {
                    if (
                        !parentRef.current?.contains(e.relatedTarget) &&
                        gridNavigationEnabled
                    ) {
                        setHasFocus(false)
                    }
                }}
                style={{ marginBottom: -offsetAdjustment || undefined }}
                onCopy={handleCopy}
                onPaste={handlePaste}
            >
                <GridStickySignal
                    data-pvds-grid-signal="left"
                    style={{
                        height: bodyHeight,
                    }}
                />
                {isDraggingEnabled && (
                    <GridDragIndicator width={columnTotalSize} />
                )}
                <GridLayout
                    role={isTreeGrid ? 'treegrid' : 'grid'}
                    aria-label={label}
                    aria-rowcount={actualRowCount}
                    aria-colcount={columnsCount}
                    aria-multiselectable={
                        selectionMode === 'multi' || isSpreadsheet
                    }
                    aria-readonly={gridCanEdit === false || undefined}
                    style={{ minWidth: columnWidth }}
                >
                    <GridHeader />
                    {!renderEmpty && (
                        <GridBody
                            role="rowgroup"
                            ref={setNodeRef}
                            style={{ height: bodyHeight }}
                        >
                            {rowVirtualItems.map((virtualItem) => {
                                const { index } = virtualItem
                                const rowId = rowIds[index]
                                if (rowId == null) {
                                    /* There is a risk that an extra row is rendered for focus, but
                            it has been removed before the range extractor is aware of it. */
                                    return null
                                }
                                return (
                                    <GridRow
                                        key={rowId}
                                        id={rowId}
                                        virtualHeight={virtualItem.size}
                                        virtualStart={virtualItem.start}
                                        virtualColumns={columnVirtualItems}
                                        stickyColumns={stickyColumnDetails}
                                        width={columnTotalSize}
                                    />
                                )
                            })}
                            {/* <DragCollisionPreview /> */}
                        </GridBody>
                    )}
                    {footerEnabled && !isTreeGrid ? (
                        <GridFooter>
                            <GridFooterRow height={rowHeight} />
                        </GridFooter>
                    ) : null}
                </GridLayout>
                <GridStickySignal
                    data-pvds-grid-signal="right"
                    style={{
                        left: columnTotalSize - 1,
                        height: bodyHeight,
                    }}
                />
                {renderEmpty && (
                    <GridEmptyContainer>{emptyContent}</GridEmptyContainer>
                )}
                <GridLiveRegion />
                <GridRangeSelection scroller={parentRef.current} />
            </ListWrap>
            <GridStickyOverlay
                totalWidth={columnWidth}
                totalHeight={bodyHeight}
                scrollingRef={parentRef}
            />
            <GridContextActionsMenu />
            <GridColumnMenu />
        </GridContainer>
    )
}
