import React, { ReactNode, isValidElement, Children } from "react";
import { ButtonIcon } from "./ButtonIcon";
import { ButtonLoading } from "./ButtonLoading";

export function useButtonSlot(children: ReactNode) {
  const iconNodes: React.ReactElement[] = [];
  const loadingNodes: React.ReactElement[] = [];

  const collectNodes = (nodes: ReactNode) => {
    Children.forEach(nodes, (child) => {
      if (!isValidElement(child)) return;

      if (child.type === ButtonIcon) {
        iconNodes.push(child);
      } else if (child.type === ButtonLoading) {
        loadingNodes.push(child);
      } else if ((child as any).props.children) {
        collectNodes((child as any).props.children);
      }
    });
  };

  collectNodes(children);

  return {
    iconNode: iconNodes.at(-1),
    loadingNode: loadingNodes.at(-1),
  };
}