import type { GridRowId } from '../types'
import type { Collection } from './types'

export function setSelection(
    renderedIds: GridRowId[],
    collection: Collection,
    selection: Set<GridRowId>,
    rowIds: GridRowId | [fromId: GridRowId, toId: GridRowId],
    selectedValue: boolean,
    matchedIds?: Set<GridRowId> | null
) {
    let ids: GridRowId[]

    if (Array.isArray(rowIds)) {
        /* Start by selecting amongst the render rows */
        const [fromId, toId] = rowIds

        const idx1 = renderedIds.indexOf(fromId)
        const idx2 = renderedIds.indexOf(toId)

        const [low, high] = idx1 < idx2 ? [idx1, idx2] : [idx2, idx1]

        ids = renderedIds.slice(low, high + 1)
    } else {
        ids = [rowIds]
    }

    const collect = (
        nodeId: GridRowId,
        memo: Set<GridRowId>,
        traverseOnlyGroups = false
    ) => {
        if (matchedIds && !matchedIds.has(nodeId)) {
            return
        }
        const meta = collection.meta.get(nodeId)
        const isGroup = meta?.type === 'group'
        if (!isGroup) {
            memo.add(nodeId)
        }
        if (!traverseOnlyGroups || (traverseOnlyGroups && isGroup)) {
            if (meta?.children?.length) {
                meta.children.forEach((childId) => collect(childId, memo))
            }
        }
    }

    const memo = new Set<GridRowId>()

    /* Collect each id, and all leaf nodes of each group */
    ids.forEach((id) => collect(id, memo, true))

    if (selectedValue) {
        return new Set([...selection, ...memo])
    } else if ([...memo].some((m) => selection.has(m))) {
        return new Set([...selection].filter((m) => !memo.has(m)))
    }

    return selection
}