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

export function flatten(
    collection: Collection,
    expanded: Set<GridRowId>,
    sortFn:
        | null
        | ((
              ids: GridRowId[],
              entities: Map<GridRowId, any>,
              meta: Map<GridRowId, GridRowMeta>
          ) => GridRowId[]),
    matchedIds?: Set<GridRowId> | null
) {
    /* Start by grabbing all the root ids, optionally filtering them if a filter is present */
    const ids = matchedIds
        ? collection.ids.filter((id) => matchedIds.has(id))
        : collection.ids

    /* Running index will serve to provide aria-rowindex that is 
       accurate regardless of rows being expanded or collapsed */
    let runningIndex = 0

    const collect = (
        nodeId: GridRowId,
        /** The ordered collection of ids, including only expanded rows */
        memo: GridRowId[],
        /** The mapping of row id to specific index, including all rows.  */
        indexes: Map<GridRowId, number>,
        /** The mapping of row id to details for aria-setsize and aria-posinset */
        setDetails: Map<GridRowId, { setSize: number; posInset: number }>,
        /** Is this level of recursion meant to still collect ids, or only indexes */
        collectIds = true
    ) => {
        const meta = collection.meta.get(nodeId)
        if (collectIds) {
            memo.push(nodeId)
        }
        indexes.set(nodeId, runningIndex)
        runningIndex++

        /* Only sort expanded, filtered children */
        if (meta?.children?.length) {
            const childIds = matchedIds
                ? meta.children.filter((id) => matchedIds.has(id))
                : meta.children
            const setSize = childIds.length

            if (expanded.has(nodeId) && collectIds) {
                const sortedChildIds = sortFn
                    ? sortFn(childIds, collection.entities, collection.meta)
                    : childIds
                /* If needed, we sort on each level as long as we are still in an expanded tree */
                sortedChildIds.forEach((childId, index) => {
                    collect(childId, memo, indexes, setDetails, true)
                    setDetails.set(childId, { setSize, posInset: index + 1 })
                })
            } else {
                /* We are in a collapsed tree, so we no longer need to sort – only collect the indexes.
                   The indexes in the collapsed section won't be accurate, but they won't be rendered so
                   that is ok. Its only the rendered items where index accuracy matters */

                childIds.forEach((childId, index) => {
                    collect(childId, memo, indexes, setDetails, false)
                    setDetails.set(childId, { setSize, posInset: index + 1 })
                })
            }
        }
    }

    const sortedRootIds = sortFn
        ? sortFn(ids, collection.entities, collection.meta)
        : ids
    const setSize = sortedRootIds.length
    const memo: GridRowId[] = []
    const indexes = new Map<GridRowId, number>()
    const setDetails = new Map<
        GridRowId,
        { setSize: number; posInset: number }
    >()

    sortedRootIds.forEach((id, index) => {
        collect(id, memo, indexes, setDetails)
        setDetails.set(id, { setSize, posInset: index + 1 })
    })

    return {
        ids: memo,
        indexes,
        setDetails,
    }
}

export function precalculateCollections(collection: Collection) {
    let rowCount = 0

    const collect = (
        nodeId: GridRowId,
        levelMemo: Map<GridRowId, number>,
        depth = 0
    ) => {
        const meta = collection.meta.get(nodeId)
        rowCount++
        levelMemo.set(nodeId, depth)

        if (meta?.children?.length) {
            meta.children.forEach((childId) =>
                collect(childId, levelMemo, depth + 1)
            )
        }
    }

    const levels = new Map<GridRowId, number>()

    collection.ids.forEach((id) => collect(id, levels))

    return {
        rowCount,
        levels,
    }
}