<?php
    /**
     * Menu cutoff walker
     * 
     * @package Blog Bee pro
     * @since 1.0.0
     */
    namespace BBE;

    use Walker_Nav_Menu;
    use Walker_Page;

    if( ! class_exists( __NAMESPACE__ . '\\Menu_Cutoff_Walker' ) ) {
        /**
         * Walker_Nav_Menu class
         */
        class Menu_Cutoff_Walker extends Walker_Nav_Menu {
            /**
             * Max number of menu items to show
             * 
             * @since 1.0.0
             */
            private $max_items;

            /**
             * Menu cutoff enabled
             * 
             * @since 1.0.0
             */
            private $cutoff_enabled;

            /**
             * Instantiate class
             * 
             * @since 1.0.0
             */
            public function __construct( $max_items = 5, $cutoff_enabled = false ) {
                $this->max_items = $max_items;
                $this->cutoff_enabled = $cutoff_enabled;
            }

            /**
             * Displays array of elements hierarchically.
             *
             * @since 1.0.0 
             * @see Walker_Nav_Menu::walk()
             * @override
             */
            public function walk( $elements, $max_depth, ...$args ) {
                if( ! $this->cutoff_enabled ) return parent::walk( $elements, $max_depth, ...$args );
                if ( empty( $elements ) ) {
                    return '';
                }

                // Normalize: ensure args[0] is always an object
                if ( isset( $args[0] ) && is_array( $args[0] ) ) {
                    $args[0] = (object) $args[0];
                } elseif ( ! isset( $args[0] ) ) {
                    $args[0] = (object) [];
                }

                // Build parent -> children map
                $children_map = [];
                foreach( $elements as $el ) {
                    $parent_id = ( int ) $el->menu_item_parent;
                    $children_map[ $parent_id ][] = $el;
                }

                $top_level     = $children_map[ 0 ] ?? [];
                $visible_top   = array_slice( $top_level, 0, $this->max_items );
                $overflow_top  = array_slice( $top_level, $this->max_items );

                $collect_branch = function( $item, &$bucket ) use ( &$collect_branch, $children_map ) {
                    $bucket[] = $item;
                    if ( isset( $children_map[ $item->ID ] ) ) {
                        foreach ( $children_map[ $item->ID ] as $child ) {
                            $collect_branch( $child, $bucket );
                        }
                    }
                };

                $visible_elements = [];
                foreach ( $visible_top as $item ) {
                    $collect_branch( $item, $visible_elements );
                }

                $overflow_elements = [];
                foreach ( $overflow_top as $item ) {
                    $collect_branch( $item, $overflow_elements );
                }

                $output = parent::walk( $visible_elements, $max_depth, ...$args );

                if ( ! empty( $overflow_elements ) ) {
                    $output .= '<li class="menu-item menu-item-has-children menu-item-more">';
                    $output .= '<a href="#">More<i class="fa-solid fa-chevron-down"></i></a>';
                    $output .= '<ul class="sub-menu">';
                    $output .= parent::walk( $overflow_elements, $max_depth, ...$args );
                    $output .= '</ul></li>';
                }

                return $output;
            }

            /**
             * Add icon in menu items that have children
             * 
             * @since 1.0.0
             * @see Walker_Nav_Menu::start_el()
             */
            public function start_el( &$output, $item, $depth = 0, $args = null, $id = 0 ) {

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

                // Detect children via class
                $has_children = in_array( 'menu-item-has-children', $classes );

                $class_names = implode( ' ', $classes );
                $output .= '<li class="' . esc_attr( $class_names ) . '">';

                // Link attributes
                $atts  = ' href="' . esc_url( $item->url ) . '"';

                $title = apply_filters( 'the_title', $item->title, $item->ID );
                if ( $has_children ) {
                    $atts .= ' aria-haspopup="true" aria-expanded="false"';
                }

                $output .= '<a' . $atts . '>';
                $output .= $title;
                // Add icon if has children
                if ( $has_children ) {
                    $output .= ' <span class="submenu-icon"><i class="fa-solid fa-chevron-down"></i></span>';
                    $output .= ' <button class="submenu-icon icon-button" type="button"><i class="fa-solid fa-chevron-down"></i></button>';
                }
                $output .= '</a>';
            }
        }
    }

    if ( ! class_exists( __NAMESPACE__ . '\\Page_Cutoff_Walker' ) ) {
        /**
         * Page_Cutoff_Walker class
         */
        class Page_Cutoff_Walker extends Walker_Page {
            /**
             * Max number of menu items to show
             * 
             * @since 1.0.0
             */
            private $max_items;

            /**
             * Menu cutoff enabled
             * 
             * @since 1.0.0
             */
            private $cutoff_enabled;

            private $children_map = [];

            /**
             * Instantiate class
             * 
             * @since 1.0.0
             */
            public function __construct( $max_items = 5, $cutoff_enabled = false ) {
                $this->max_items = $max_items;
                $this->cutoff_enabled = $cutoff_enabled;
            }

            /**
             * Walk through elements and limit top-level items
             */
            public function walk( $elements, $max_depth, ...$args ) {
                if ( empty( $elements ) ) {
                    return '';
                }

                // Build parent -> children map
                foreach( $elements as $el ) {
                    $parent_id = ( int ) $el->post_parent;
                    $this->children_map[ $parent_id ][] = $el;
                }

                if( ! $this->cutoff_enabled ) return parent::walk( $elements, $max_depth, ...$args );

                $children_map  = $this->children_map;
                $top_level     = $children_map[ 0 ] ?? [];
                $visible_top   = array_slice( $top_level, 0, $this->max_items );
                $overflow_top  = array_slice( $top_level, $this->max_items );

                $collect_branch = function( $item, &$bucket ) use ( &$collect_branch, $children_map ) {
                    $bucket[] = $item;
                    if ( isset( $children_map[ $item->ID ] ) ) {
                        foreach ( $children_map[ $item->ID ] as $child ) {
                            $collect_branch( $child, $bucket );
                        }
                    }
                };

                $visible_elements = [];
                foreach ( $visible_top as $item ) {
                    $collect_branch( $item, $visible_elements );
                }

                $overflow_elements = [];
                foreach ( $overflow_top as $item ) {
                    $collect_branch( $item, $overflow_elements );
                }

                $output = parent::walk( $visible_elements, $max_depth, ...$args );

                if ( ! empty( $overflow_elements ) ) {
                    $output .= '<li class="menu-item menu-item-has-children menu-item-more">';
                    $output .= '<a href="#">More<i class="fa-solid fa-chevron-down"></i></a>';
                    $output .= '<ul class="sub-menu">';
                    $output .= parent::walk( $overflow_elements, $max_depth, ...$args );
                    $output .= '</ul></li>';
                }

                return $output;
            }

            /**
             * Add icon in menu items that have children
             * 
             * @since 1.0.0
             * @see Walker_Nav_Menu::start_el()
             */
            public function start_el(&$output, $item, $depth = 0, $args = null, $id = 0) {

                $classes = isset($item->classes) ? (array) $item->classes : [];

                // Correct child detection for Walker_Page
                $has_children = !empty($this->children_map[$item->ID]);

                // Ensure correct base class is always present
                $classes[] = 'page_item';
                $classes[] = 'page-item-' . $item->ID;

                if ($has_children) {
                    $classes[] = 'page_item_has_children';
                }

                $class_names = implode(' ', array_filter($classes));

                $output .= '<li class="' . esc_attr($class_names) . '">';

                $atts = ' href="' . esc_url(get_permalink($item->ID)) . '"';

                $title = apply_filters('the_title', $item->post_title, $item->ID);
                if ( $has_children ) {
                    $atts .= ' aria-haspopup="true" aria-expanded="false"';
                }

                $output .= '<a' . $atts . '>';

                $output .= $title;

                if( $has_children ) {
                    $output .= ' <span class="submenu-icon"><i class="fa-solid fa-chevron-down"></i></span>';
                    $output .= ' <button class="submenu-icon icon-button" type="button"><i class="fa-solid fa-chevron-down"></i></button>';
                }

                $output .= '</a>';
            }
        }
    }