Issues with Custom Walker_Nav_Menu for Columns in WordPress Navigation Menus

I’m trying to create a custom navigation menu in WordPress that uses invisible columns as containers for second-level menu items. I’m using a custom class that extends Walker_Nav_Menu, but I’m having trouble properly nesting the menu items within these columns.

Desired Menu Structure

Here’s how I want the structure of my menu to look:

<div class="children-wrapper">
    <div class="column-container">
        <li id="menu-item-98468" class="menu-column column-1 menu-item menu-item-type-custom menu-item-object-custom menu-item-98468">
            <a href="#">COLUMN-1</a>
            <ul class="sub-menu">
                <!-- Second-level items -->
            </ul>
        </li>
        <li id="menu-item-98460" class="menu-column column-2 menu-item menu-item-type-custom menu-item-object-custom menu-item-98460">
            <a href="#">COLUMN-2</a>
            <ul class="sub-menu">
                <!-- Second-level items -->
            </ul>
        </li>
    </div>
</div>

Current Menu Structure

<?php
class CustomWalker extends Walker_Nav_Menu {
    private $column_index = 0;
    private $in_column = false;

    function start_lvl(&$output, $depth = 0, $args = array()) {
        $indent = str_repeat("t", $depth);
        $n = "n";

        if ($depth === 0) {
            $output .= "{$n}{$indent}<div class="children-wrapper">{$n}";
            $output .= "{$indent}<div class="column-container">{$n}";
        } else {
            $classes = array('sub-menu');
            $class_names = join(' ', apply_filters('nav_menu_submenu_css_class', $classes, $args, $depth));
            $class_names = $class_names ? ' class="' . esc_attr($class_names) . '"' : '';
            $output .= "{$n}{$indent}<ul$class_names>{$n}";
        }
    }

    function end_lvl(&$output, $depth = 0, $args = array()) {
        $indent = str_repeat("t", $depth);
        $n = "n";

        if ($depth === 0) {
            $output .= "{$indent}</div><!-- .column-container -->{$n}";
            $output .= "{$indent}</div><!-- .children-wrapper -->{$n}";
        } else {
            $output .= "{$indent}</ul>{$n}";
        }
    }

    function start_el(&$output, $item, $depth = 0, $args = array(), $id = 0) {
        error_log("Depth: " . $depth . ", Item ID: " . $item->ID . ", Title: " . $item->title);
        $indent = ($depth) ? str_repeat("t", $depth) : '';
        $n = "n";

        $classes = empty($item->classes) ? array() : (array) $item->classes;

        // Check if this item is a column
        if (in_array('column-1', $classes) || in_array('column-2', $classes) || in_array('column-3', $classes)) {
            $this->in_column = true;
            $this->column_index++;
            $output .= "{$indent}<div class="menu-column column-{$this->column_index}">{$n}";
            $output .= "{$indent}<ul class="sub-menu">{$n}";
            return; // Don't output the column item itself
        }

        $this->output_item($output, $item, $depth, $args, $id);
    }

    function end_el(&$output, $item, $depth = 0, $args = array()) {
        $n = "n";

        $classes = empty($item->classes) ? array() : (array) $item->classes;

        if (in_array('column-1', $classes) || in_array('column-2', $classes) || in_array('column-3', $classes)) {
            $output .= "</ul>{$n}";
            $output .= "</div><!-- .menu-column -->{$n}";
            $this->in_column = false;
        } else {
            $output .= "</li>{$n}";
        }
    }

    private function output_item(&$output, $item, $depth, $args, $id) {
        $indent = ($depth) ? str_repeat("t", $depth) : '';
        $n = "n";

        $classes = empty($item->classes) ? array() : (array) $item->classes;
        $classes[] = 'menu-item-' . $item->ID;
        $class_names = join(' ', apply_filters('nav_menu_css_class', array_filter($classes), $item, $args, $depth));
        $class_names = $class_names ? ' class="' . esc_attr($class_names) . '"' : '';

        $id = apply_filters('nav_menu_item_id', 'menu-item-' . $item->ID, $item, $args, $depth);
        $id = $id ? ' id="' . esc_attr($id) . '"' : '';

        $output .= $indent . '<li' . $id . $class_names .'>';

        $atts = array();
        $atts['title'] = !empty($item->attr_title) ? $item->attr_title : '';
        $atts['target'] = !empty($item->target) ? $item->target : '';
        $atts['rel'] = !empty($item->xfn) ? $item->xfn : '';
        $atts['href'] = !empty($item->url) ? $item->url : '';

        $atts = apply_filters('nav_menu_link_attributes', $atts, $item, $args, $depth);

        $attributes = '';
        foreach ($atts as $attr => $value) {
            if (!empty($value)) {
                $value = ('href' === $attr) ? esc_url($value) : esc_attr($value);
                $attributes .= ' ' . $attr . '="' . $value . '"';
            }
        }

        $title = apply_filters('the_title', $item->title, $item->ID);

        $item_output = $args->before;
        $item_output .= '<a'. $attributes .'>';
        $item_output .= $args->link_before . $title . $args->link_after;
        $item_output .= '</a>';
        $item_output .= $args->after;

        $output .= apply_filters('walker_nav_menu_start_el', $item_output, $item, $depth, $args);
    }
}
?>

Current Problem

The column items are being generated outside the column-container, and the second-level items are not being properly nested within these columns. I also think that the depthvariable isn’t working as it should be.
I’ll attach an image of what I’d like the menu structure to be

What I’ve Tried

I’ve tried modifying the start_el and end_el methods to properly nest the column items within the column-container, but the columns are still being generated outside of it. Any suggestions on how to fix this issue and properly nest the columns within the column-container?

Thank you in advance for your help!