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

const NO_META = {}

export function filter(
    collection: Collection,
    filterFn: (row: GridRowData, meta: GridRowMeta) => boolean
): {
    actualMatches: Set<GridRowId>
    necessaryMatches: Set<GridRowId>
} {
    /**
     * In normal filtering mode, we match any parent or child.
     *  - If a parent matches, all of its children match.
     *  - If a child matches, all of its parents match
     *  - If a parent doesn't match, and none of its descendants match, it is removed.
     */

    /* Used when a parent matches, and we just need to get all the children */
    const collect = (
        nodeId: GridRowId,
        memo: Set<GridRowId>,
        memoWithParents: Set<GridRowId>
    ) => {
        const meta = collection.meta.get(nodeId)
        memo.add(nodeId)
        memoWithParents.add(nodeId)
        if (meta?.children?.length) {
            meta.children.forEach((childId) =>
                collect(childId, memo, memoWithParents)
            )
        }
    }

    /* Used when matching on each level */
    const collectIfMatching = (
        nodeId: GridRowId,
        memo: Set<GridRowId>,
        memoWithParents: Set<GridRowId>
    ) => {
        const node = collection.entities.get(nodeId)
        const meta = collection.meta.get(nodeId)

        const matches = node && filterFn(node, meta || NO_META)
        if (matches) {
            memo.add(nodeId)
            memoWithParents.add(nodeId)
        }
        // The parent matches, collect all the children
        if (meta?.children?.length) {
            if (matches) {
                meta.children.forEach((childId) =>
                    collect(childId, memo, memoWithParents)
                )
            } else {
                const matchingChildren = new Set<GridRowId>()
                const matchingChildrenWithParents = new Set<GridRowId>()
                meta.children.forEach((childId) =>
                    collectIfMatching(
                        childId,
                        matchingChildren,
                        matchingChildrenWithParents
                    )
                )
                if (matchingChildren.size) {
                    ;[...matchingChildren].forEach((childId) => {
                        memo.add(childId)
                    })
                }
                if (matchingChildrenWithParents.size) {
                    memoWithParents.add(nodeId)
                    ;[...matchingChildrenWithParents].forEach((childId) => {
                        memoWithParents.add(childId)
                    })
                }
            }
        }
    }

    /* These are ids that either exactly matched the filter, or are descendants of
           and exact match and thus match */
    const actualMatches = new Set<GridRowId>()

    /* These are the same as actualMatches, with the addition of any ancestor
           hierarchy to properly render a matched item */
    const necessaryMatches = new Set<GridRowId>()

    collection.ids.forEach((id) =>
        collectIfMatching(id, actualMatches, necessaryMatches)
    )

    return {
        actualMatches,
        necessaryMatches: necessaryMatches,
    }
}