[dyad] Add collapse/expand all to headlines - wrote 1 file(s)

src/components/headline-tree.tsx

@@ -6,24 +6,48 @@ import type { HeadlineNode } from "@/app/actions";
 import { KeywordHighlighter } from "./keyword-highlighter";
 import { ChevronDown } from "lucide-react";
 import { cn } from "@/lib/utils";
+import { Button } from "@/components/ui/button";
 
 interface HeadlineTreeProps {
   headlines: HeadlineNode[];
   keyword?: string | null;
 }
 
+const getCollapsiblePaths = (
+  nodes: HeadlineNode[],
+  parentPath = ""
+): string[] => {
+  let paths: string[] = [];
+  nodes.forEach((node, index) => {
+    const currentPath = parentPath ? `${parentPath}-${index}` : `${index}`;
+    if (node.level === 2 && node.children && node.children.length > 0) {
+      paths.push(currentPath);
+    }
+    if (node.children && node.children.length > 0) {
+      paths = paths.concat(getCollapsiblePaths(node.children, currentPath));
+    }
+  });
+  return paths;
+};
+
 const HeadlineNodeDisplay = ({
   node,
   level,
+  path,
   keyword,
+  collapsedStates,
+  onToggle,
 }: {
   node: HeadlineNode;
   level: number;
+  path: string;
   keyword?: string | null;
+  collapsedStates: Record<string, boolean>;
+  onToggle: (path: string) => void;
 }) => {
   const isCollapsible =
     node.level === 2 && node.children && node.children.length > 0;
-  const [isCollapsed, setIsCollapsed] = useState(true);
+  const isCollapsed = collapsedStates[path];
 
   return (
     <>
@@ -33,7 +57,7 @@ const HeadlineNodeDisplay = ({
           isCollapsible && "cursor-pointer hover:bg-muted/50"
         )}
         style={{ paddingLeft: `${16 + level * 24}px` }}
-        onClick={() => isCollapsible && setIsCollapsed(!isCollapsed)}
+        onClick={() => isCollapsible && onToggle(path)}
       >
         <Badge
           variant="secondary"
@@ -63,7 +87,10 @@ const HeadlineNodeDisplay = ({
             key={index}
             node={child}
             level={level + 1}
+            path={`${path}-${index}`}
             keyword={keyword}
+            collapsedStates={collapsedStates}
+            onToggle={onToggle}
           />
         ))}
     </>
@@ -71,15 +98,64 @@ const HeadlineNodeDisplay = ({
 };
 
 export function HeadlineTree({ headlines, keyword }: HeadlineTreeProps) {
+  const [collapsedStates, setCollapsedStates] = useState<
+    Record<string, boolean>
+  >(() => {
+    const initialStates: Record<string, boolean> = {};
+    const paths = getCollapsiblePaths(headlines);
+    paths.forEach((path) => {
+      initialStates[path] = true;
+    });
+    return initialStates;
+  });
+
+  const handleToggle = (path: string) => {
+    setCollapsedStates((prev) => ({
+      ...prev,
+      [path]: !prev[path],
+    }));
+  };
+
+  const collapseAll = () => {
+    const newStates: Record<string, boolean> = {};
+    Object.keys(collapsedStates).forEach((path) => {
+      newStates[path] = true;
+    });
+    setCollapsedStates(newStates);
+  };
+
+  const expandAll = () => {
+    const newStates: Record<string, boolean> = {};
+    Object.keys(collapsedStates).forEach((path) => {
+      newStates[path] = false;
+    });
+    setCollapsedStates(newStates);
+  };
+
+  const hasCollapsibleNodes = Object.keys(collapsedStates).length > 0;
+
   return (
     <div className="w-full rounded-lg overflow-hidden">
+      {hasCollapsibleNodes && (
+        <div className="flex items-center gap-2 p-4 border-b bg-muted/50">
+          <Button variant="outline" size="sm" onClick={expandAll}>
+            Expand All
+          </Button>
+          <Button variant="outline" size="sm" onClick={collapseAll}>
+            Collapse All
+          </Button>
+        </div>
+      )}
       <div className="-mt-px">
         {headlines.map((headline, index) => (
           <HeadlineNodeDisplay
             key={index}
             node={headline}
             level={0}
+            path={`${index}`}
             keyword={keyword}
+            collapsedStates={collapsedStates}
+            onToggle={handleToggle}
           />
         ))}
       </div>