From 242ffd2e8d6eb0d8ae26d533d1209ebca1ef8fb7 Mon Sep 17 00:00:00 2001
From: "nida.ghuman" <nidaqg@gmail.com>
Date: Mon, 23 Dec 2024 16:35:18 -0500
Subject: [PATCH] wrapped column logic

---
 .../pb_advanced_table/table_header.html.erb   |  7 ++--
 .../pb_advanced_table/table_header.rb         | 33 ++++++++++++++++---
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/playbook/app/pb_kits/playbook/pb_advanced_table/table_header.html.erb b/playbook/app/pb_kits/playbook/pb_advanced_table/table_header.html.erb
index 8542074ead..02253179fe 100644
--- a/playbook/app/pb_kits/playbook/pb_advanced_table/table_header.html.erb
+++ b/playbook/app/pb_kits/playbook/pb_advanced_table/table_header.html.erb
@@ -2,12 +2,10 @@
   <% object.header_rows.each_with_index do |header_row, row_index| %>
     <%= pb_rails("table/table_row", props: { tag: "div" }) do %>
       <% header_row.each_with_index do |cell, cell_index| %>
-
         <% header_id = cell[:accessor].present? ? cell[:accessor] : "header_#{row_index}_#{cell_index}" %>
-
         <%= pb_rails("table/table_header", props: { tag: "div", id: header_id, classname: object.th_classname, sort_menu: cell[:accessor] ? cell[:sort_menu] : nil }) do %>
-                <%= pb_rails("flex", props:{ align: "center", justify: row_index.zero? ? "start" : "end", text_align: "end" }) do %>
-                    <% if cell_index.zero? && (object.enable_toggle_expansion == "header" || object.enable_toggle_expansion == "all") %>
+                <%= pb_rails("flex", props:{ align: "center", justify: cell_index.zero? ? "start" : "end", text_align: "end" }) do %>
+                    <% if cell_index.zero? && (object.enable_toggle_expansion == "header" || object.enable_toggle_expansion == "all") && row_index === header_rows.size - 1 %>
                         <button
                             class="gray-icon toggle-all-icon"
                             onclick="expandAllRows(this); event.preventDefault();">
@@ -17,7 +15,6 @@
                     <%= cell[:label] %>
                 <% end %>
         <% end %>
-
       <% end %>
     <% end %>
   <% end %>
diff --git a/playbook/app/pb_kits/playbook/pb_advanced_table/table_header.rb b/playbook/app/pb_kits/playbook/pb_advanced_table/table_header.rb
index 39245498d8..fe101b5f82 100644
--- a/playbook/app/pb_kits/playbook/pb_advanced_table/table_header.rb
+++ b/playbook/app/pb_kits/playbook/pb_advanced_table/table_header.rb
@@ -18,13 +18,12 @@ def th_classname
       end
 
       def header_rows
-        rows = []
-        max_depth = compute_max_depth(column_definitions)
+        wrapped_columns = wrap_leaf_columns(column_definitions)
 
+        rows = []
+        max_depth = compute_max_depth(wrapped_columns)
         max_depth.times { rows << [] }
-
-        process_columns(column_definitions, rows, 0, max_depth)
-
+        process_columns(wrapped_columns, rows, 0, max_depth)
         rows
       end
 
@@ -63,6 +62,30 @@ def compute_leaf_columns(columns)
           col[:columns] ? sum + compute_leaf_columns(col[:columns]) : sum + 1
         end
       end
+
+      def wrap_leaf_columns(column_definitions)
+        has_grouped_headers = column_definitions.any? { |col| col.key?(:columns) }
+
+        if has_grouped_headers
+          column_definitions.map do |col|
+            if col.key?(:columns)
+              {
+                label: col[:label],
+                columns: wrap_leaf_columns(col[:columns]),
+              }
+            elsif col.key?(:accessor)
+              {
+                label: "",
+                columns: [col],
+              }
+            else
+              col
+            end
+          end
+        else
+          column_definitions
+        end
+      end
     end
   end
 end