children.ts

import React, {
  Children,
  cloneElement,
  HTMLAttributes,
  isValidElement,
} from "react"

interface Schema {
  [key: string]:
    | HTMLAttributes<any>
    | ((
        child: React.ReactElement,
        index: number
      ) => React.ReactElement | HTMLAttributes<any>)
}

export const remapChildren = (
  children: React.ReactNode,
  schema: Schema
): React.ReactNode => {
  return Children.map(children, (child, index) => {
    if (!isValidElement(child)) {
      return child
    }

    if (child.props.children) {
      child = cloneElement(child, {
        // @ts-ignore
        children: remapChildren(child.props.children, schema),
      })
    }

    for (const [key, value] of Object.entries(schema)) {
      if (
        typeof child !== "string" &&
        typeof child !== "number" &&
        typeof child.type !== "string"
      ) {
        if (child.type.name === key) {
          if (typeof value === "function") {
            const result = value(child, index)
            if (React.isValidElement(result)) {
              return result
            } else {
              return cloneElement(child, result as HTMLAttributes<any>)
            }
          } else {
            return cloneElement(child, value)
          }
        }
      }
    }

    return child
  })
}