import type { DataEntity } from '../../vc-tree/interface'
import type { InternalFieldName } from '../TreeSelect'
import type { Key, RawValueType } from '../interface'
import { isCheckDisabled } from './valueUtil'

export const SHOW_ALL = 'SHOW_ALL'
export const SHOW_PARENT = 'SHOW_PARENT'
export const SHOW_CHILD = 'SHOW_CHILD'

export type CheckedStrategy = typeof SHOW_ALL | typeof SHOW_PARENT | typeof SHOW_CHILD

export function formatStrategyValues(
  values: Key[],
  strategy: CheckedStrategy,
  keyEntities: Record<Key, DataEntity>,
  fieldNames: InternalFieldName,
): RawValueType[] {
  const valueSet = new Set(values)

  if (strategy === SHOW_CHILD) {
    return values.filter((key) => {
      const entity = keyEntities[key]

      if (
        entity &&
        entity.children &&
        entity.children.some(({ node }) => valueSet.has(node[fieldNames.value])) &&
        entity.children.every(
          ({ node }) => isCheckDisabled(node) || valueSet.has(node[fieldNames.value]),
        )
      ) {
        return false
      }
      return true
    })
  }
  if (strategy === SHOW_PARENT) {
    return values.filter((key) => {
      const entity = keyEntities[key]
      const parent = entity ? entity.parent : null
      if (parent && !isCheckDisabled(parent.node) && valueSet.has(parent.key)) {
        return false
      }
      return true
    })
  }
  return values
}
