Checkbox Tree

"use client";

import { CheckboxTree } from "@/registry/origin-ui/checkbox-tree";
import { Checkbox } from "@/components/ui/checkbox";
import { Label } from "@/components/ui/label";

const tree = {
  id: "root",
  label: "All Categories",
  children: [
    {
      id: "frontend",
      label: "Frontend",
      children: [
        { id: "react", label: "React", defaultChecked: true },
        { id: "vue", label: "Vue" },
        { id: "svelte", label: "Svelte", defaultChecked: true },
      ],
    },
    {
      id: "backend",
      label: "Backend",
      children: [
        { id: "node", label: "Node.js", defaultChecked: true },
        { id: "python", label: "Python" },
        { id: "go", label: "Go" },
      ],
    },
    {
      id: "database",
      label: "Database",
      children: [
        { id: "postgres", label: "PostgreSQL" },
        { id: "mongo", label: "MongoDB" },
      ],
    },
  ],
};

export default function Demo() {
  return (
    <div className="flex items-center justify-center min-h-[300px] p-8">
      <div className="w-64 space-y-1">
        <CheckboxTree
          tree={tree}
          renderNode={({ node, isChecked, onCheckedChange, children }) => (
            <div key={node.id} className="space-y-1">
              <div className="flex items-center gap-2 py-0.5">
                <Checkbox
                  id={node.id}
                  checked={isChecked === true ? true : isChecked === "indeterminate" ? "indeterminate" : false}
                  onCheckedChange={onCheckedChange}
                />
                <Label htmlFor={node.id} className="cursor-pointer font-normal">
                  {node.label}
                </Label>
              </div>
              {children && (
                <div className="pl-6 space-y-1 border-l border-border ml-2">
                  {children}
                </div>
              )}
            </div>
          )}
        />
      </div>
    </div>
  );
}

/**
 * IMPORTANT: This component was built for demo purposes only and has not been tested in production.
 * It serves as a proof of concept for a checkbox tree implementation.
 * If you're interested in collaborating to create a more robust, production-ready
 * headless component, your contributions are welcome!
 */

"use client";

import type React from "react";
import { useCallback, useMemo, useState } from "react";

interface TreeNode {
  id: string;
  label: string;
  defaultChecked?: boolean;
  children?: TreeNode[];
}

function useCheckboxTree(initialTree: TreeNode) {
  const initialCheckedNodes = useMemo(() => {
    const checkedSet = new Set<string>();
    const initializeCheckedNodes = (node: TreeNode) => {
      if (node.defaultChecked) {
        checkedSet.add(node.id);
      }
      node.children?.forEach(initializeCheckedNodes);
    };
    initializeCheckedNodes(initialTree);
    return checkedSet;
  }, [initialTree]);

  const [checkedNodes, setCheckedNodes] =
    useState<Set<string>>(initialCheckedNodes);

  const isChecked = useCallback(
    (node: TreeNode): boolean | "indeterminate" => {
      if (!node.children) {
        return checkedNodes.has(node.id);
      }

      const childrenChecked = node.children.map((child) => isChecked(child));
      if (childrenChecked.every((status) => status === true)) {
        return true;
      }
      if (
        childrenChecked.some(
          (status) => status === true || status === "indeterminate",
        )
      ) {
        return "indeterminate";
      }
      return false;
    },
    [checkedNodes],
  );

  const handleCheck = useCallback(
    (node: TreeNode) => {
      const newCheckedNodes = new Set(checkedNodes);

      const toggleNode = (n: TreeNode, check: boolean) => {
        if (check) {
          newCheckedNodes.add(n.id);
        } else {
          newCheckedNodes.delete(n.id);
        }
        for (const child of n.children ?? []) {
          toggleNode(child, check);
        }
      };

      const currentStatus = isChecked(node);
      const newCheck = currentStatus !== true;

      toggleNode(node, newCheck);
      setCheckedNodes(newCheckedNodes);
    },
    [checkedNodes, isChecked],
  );

  return { handleCheck, isChecked };
}

interface CheckboxTreeProps {
  tree: TreeNode;
  renderNode: (props: {
    node: TreeNode;
    isChecked: boolean | "indeterminate";
    onCheckedChange: () => void;
    children: React.ReactNode;
  }) => React.ReactNode;
}

export function CheckboxTree({ tree, renderNode }: CheckboxTreeProps) {
  const { isChecked, handleCheck } = useCheckboxTree(tree);

  const renderTreeNode = (node: TreeNode): React.ReactNode => {
    const children = node.children?.map(renderTreeNode);

    return renderNode({
      children,
      isChecked: isChecked(node),
      node,
      onCheckedChange: () => handleCheck(node),
    });
  };

  return renderTreeNode(tree);
}