Untitled

#include <memory>
#include <wayfire/plugin.hpp>
#include <wayfire/per-output-plugin.hpp>
#include <wayfire/output.hpp>
#include <wayfire/core.hpp>
#include <wayfire/view.hpp>
#include <wayfire/toplevel-view.hpp>
#include <wayfire/seat.hpp>
#include <wayfire/signal-definitions.hpp>
#include <wayfire/util/log.hpp>
#include <wayfire/view-transform.hpp>
#include <wayfire/render-manager.hpp>
#include <wayfire/scene-operations.hpp>
#include <wayfire/opengl.hpp>

// Helper to check for intersection
bool boxes_intersect(const wlr_box& a, const wlr_box& b)
{
    return !(b.x >= a.x + a.width || a.x >= b.x + b.width ||
             b.y >= a.y + a.height || a.y >= b.y + b.height);
}

// Helper to check if point is inside geometry
bool point_inside(const wf::geometry_t& geom, int x, int y)
{
    return x >= geom.x && x < geom.x + geom.width &&
           y >= geom.y && y < geom.y + geom.height;
}

// Preview overlay data
struct preview_overlay_t
{
    wayfire_view target_view;
    wf::geometry_t position;
    wf::dimensions_t size;
};

// Represents a single grouped window
struct group_member_t
{
    wayfire_view view;
    std::shared_ptr<wf::scene::view_2d_transformer_t> transformer;
    std::unique_ptr<preview_overlay_t> preview_overlay;
    bool is_hidden;

    group_member_t() : is_hidden(false) {}
    group_member_t(const group_member_t&) = delete;
    group_member_t& operator=(const group_member_t&) = delete;
    group_member_t(group_member_t&&) = default;
    group_member_t& operator=(group_member_t&&) = default;
};

// Represents a full group of windows
struct window_group_t
{
    std::vector<group_member_t> members;
};

class wayfire_group : public wf::per_output_plugin_instance_t
{
private:
    // Configuration options
    wf::option_wrapper_t<int> preview_size{"group/preview_size"};
    wf::option_wrapper_t<int> preview_margin{"group/preview_margin"};

    // A map to quickly find which group a view belongs to
    std::map<wayfire_view, std::shared_ptr<window_group_t>> view_to_group_map;

    // Track the last known positions for detecting movement
    std::map<wayfire_view, wf::geometry_t> last_positions;

    // Keybinding to create/destroy groups
    wf::option_wrapper_t<wf::keybinding_t> group_key{"group/toggle"};

    // Button binding for click detection
    wf::option_wrapper_t<wf::buttonbinding_t> click_binding{"group/click"};

    // Transformer name for group movements
    std::string group_transformer_name = "group-transform";

    // Effect hook for rendering overlays on top
    wf::effect_hook_t overlay_effect = [=] ()
    {
        render_preview_overlays();
    };

    // Key callback for the keybinding
    wf::key_callback group_key_callback = [=] (const wf::keybinding_t& key) -> bool
    {
        auto focused_view = wf::get_core().seat->get_active_view();
        if (!focused_view)
        {
            return false;
        }

        if (view_to_group_map.count(focused_view))
        {
            ungroup_view(focused_view);
        } else
        {
            create_group_from_focused(focused_view);
        }

        return true;
    };

    // Button callback for click detection
    wf::button_callback click_callback = [=] (const wf::buttonbinding_t& button) -> bool
    {
        auto cursor = wf::get_core().get_cursor_position();
        return handle_preview_click(cursor.x, cursor.y);
    };

    // Handle clicks on preview overlays
    bool handle_preview_click(int x, int y)
    {
        for (const auto& [view, group] : view_to_group_map)
        {
            for (auto& member : group->members)
            {
                if (member.is_hidden && member.preview_overlay)
                {
                    auto& overlay = *member.preview_overlay;
                    wf::geometry_t click_area{
                        overlay.position.x,
                        overlay.position.y,
                        overlay.size.width,
                        overlay.size.height
                    };

                    if (point_inside(click_area, x, y))
                    {
                        switch_to_window(member, group);
                        return true; // Consume the click
                    }
                }
            }
        }
        return false; // Let the click pass through
    }

    // Switch to a specific window in the group
    void switch_to_window(group_member_t& target_member, std::shared_ptr<window_group_t> group)
    {
        if (!target_member.view || !target_member.is_hidden) return;

        LOGD("GROUP: Switching to window ", target_member.view->get_title());

        // Find current visible member (anchor)
        group_member_t* current_anchor = nullptr;
        for (auto& member : group->members)
        {
            if (!member.is_hidden)
            {
                current_anchor = &member;
                break;
            }
        }

        if (!current_anchor) return;

        // Hide current anchor
        if (!current_anchor->transformer)
        {
            // Create transformer for the current anchor
            current_anchor->transformer = std::make_shared<wf::scene::view_2d_transformer_t>(current_anchor->view);
            current_anchor->view->get_transformed_node()->add_transformer(current_anchor->transformer,
                wf::TRANSFORMER_2D, group_transformer_name);
        }
        hide_window(*current_anchor);
        create_preview_overlay(*current_anchor, target_member.view);

        // Show target member
        show_window(target_member);
        destroy_preview_overlay(target_member);

        // Focus the newly visible window
        wf::get_core().seat->focus_view(target_member.view);

        // Update preview positions
        update_preview_positions(group);

        LOGD("GROUP: Window switch complete");
    }

    // Create a preview overlay
    void create_preview_overlay(group_member_t& member, wayfire_view anchor_view)
    {
        if (!member.view || member.preview_overlay) return;

        auto toplevel = wf::toplevel_cast(member.view);
        auto anchor_toplevel = wf::toplevel_cast(anchor_view);
        if (!toplevel || !anchor_toplevel) return;

        member.preview_overlay = std::make_unique<preview_overlay_t>();
        auto& overlay = *member.preview_overlay;

        overlay.target_view = member.view;

        // Calculate preview size maintaining aspect ratio
        auto view_geom = toplevel->get_geometry();
        int target_size = preview_size;

        if (view_geom.width > view_geom.height)
        {
            overlay.size.width = target_size;
            overlay.size.height = (view_geom.height * target_size) / view_geom.width;
        }
        else
        {
            overlay.size.height = target_size;
            overlay.size.width = (view_geom.width * target_size) / view_geom.height;
        }

        // Position in top-right of anchor window
        auto anchor_geom = anchor_toplevel->get_geometry();
        overlay.position.x = anchor_geom.x + anchor_geom.width - overlay.size.width - preview_margin;
        overlay.position.y = anchor_geom.y + preview_margin;

        LOGD("GROUP: Created preview overlay for ", member.view->get_title());
    }

    // Destroy preview overlay
    void destroy_preview_overlay(group_member_t& member)
    {
        if (!member.preview_overlay) return;
        member.preview_overlay.reset();
        LOGD("GROUP: Destroyed preview overlay for ", member.view->get_title());
    }

    // Hide a window by scaling it to zero
    void hide_window(group_member_t& member)
    {
        if (!member.transformer || member.is_hidden) return;
        
        // --- START OF FIX ---
        // Damage the view's current area BEFORE hiding it by scaling to zero.
        // This is crucial to notify the compositor that the area previously occupied
        // by this window needs to be redrawn, thus clearing its old contents and
        // preventing visual artifacts (ghost windows).
        member.view->damage();
        // --- END OF FIX ---

        // Scale to zero to hide the window
        member.transformer->scale_x = 0.0f;
        member.transformer->scale_y = 0.0f;
        member.is_hidden = true;
        member.view->damage(); // Damage after is also good practice.

        LOGD("GROUP: Hidden window ", member.view->get_title());
    }

    // Show a window by restoring its scale
    void show_window(group_member_t& member)
    {
        if (!member.transformer || !member.is_hidden) return;

        // Restore normal scale
        member.transformer->scale_x = 1.0f;
        member.transformer->scale_y = 1.0f;
        member.is_hidden = false;
        member.view->damage();

        LOGD("GROUP: Shown window ", member.view->get_title());
    }

    // Update preview positions
    void update_preview_positions(std::shared_ptr<window_group_t> group)
    {
        if (group->members.empty()) return;

        // Find the anchor (the single visible member in the group)
        group_member_t* anchor = nullptr;
        for (auto& member : group->members)
        {
            if (!member.is_hidden)
            {
                anchor = &member;
                break;
            }
        }

        if (!anchor) return;

        auto anchor_toplevel = wf::toplevel_cast(anchor->view);
        if (!anchor_toplevel) return;

        // Calculate the actual on-screen geometry of the anchor window.
        // get_geometry() returns the base geometry. If a transformer is active,
        // we must account for its translation to find the correct visual position.
        wf::geometry_t anchor_geom = anchor_toplevel->get_geometry();
        if (anchor->transformer)
        {
            anchor_geom.x += anchor->transformer->translation_x;
            anchor_geom.y += anchor->transformer->translation_y;
        }

        int current_y = anchor_geom.y + preview_margin;

        // Update positions for all preview overlays
        for (auto& member : group->members)
        {
            // Only update positions for hidden members that have a preview
            if (member.is_hidden && member.preview_overlay)
            {
                auto& overlay = *member.preview_overlay;
                overlay.position.x = anchor_geom.x + anchor_geom.width - overlay.size.width - preview_margin;
                overlay.position.y = current_y;

                current_y += overlay.size.height + preview_margin;
            }
        }

        // Schedule a redraw to show updated overlays
        this->output->render->schedule_redraw();
    }

    // Render all preview overlays on top of everything
    void render_preview_overlays()
    {
        for (const auto& [view, group] : view_to_group_map)
        {
            render_group_overlays(group);
        }
    }

    // Render overlays for a specific group
    void render_group_overlays(std::shared_ptr<window_group_t> group)
    {
        if (!group || group->members.empty()) return;

        // Find anchor to check if we should render on this output
        group_member_t* anchor = nullptr;
        for (auto& member : group->members)
        {
            if (!member.transformer || !member.is_hidden)
            {
                anchor = &member;
                break;
            }
        }

        if (!anchor || anchor->view->get_output() != this->output) return;

        // Render previews for hidden group members
        for (auto& member : group->members)
        {
            if (&member != anchor && member.is_hidden && member.preview_overlay)
            {
                render_single_overlay(member);
            }
        }
    }

    // Render a single preview overlay
    void render_single_overlay(const group_member_t& member)
    {
        if (!member.preview_overlay || !member.view) return;

        auto& overlay = *member.preview_overlay;

        // Take snapshot of the view
        wf::auxilliary_buffer_t snapshot;
        member.view->take_snapshot(snapshot);

        auto snapshot_size = snapshot.get_size();
        if (snapshot_size.width > 0 && snapshot_size.height > 0)
        {
            // Calculate scaling to fit preview size
            double scale_x = double(overlay.size.width) / snapshot_size.width;
            double scale_y = double(overlay.size.height) / snapshot_size.height;
            double scale = std::min(scale_x, scale_y);

            int scaled_width = snapshot_size.width * scale;
            int scaled_height = snapshot_size.height * scale;

            // Center the scaled preview
            int offset_x = (overlay.size.width - scaled_width) / 2;
            int offset_y = (overlay.size.height - scaled_height) / 2;

            wf::geometry_t dest_geometry{
                overlay.position.x + offset_x,
                overlay.position.y + offset_y,
                scaled_width,
                scaled_height
            };

            wlr_fbox src_box{0, 0, float(snapshot_size.width), float(snapshot_size.height)};

            // Get the current render target and render the preview
            auto target_fb = this->output->render->get_target_framebuffer();
            target_fb.blit(snapshot, src_box, dest_geometry);
            snapshot.free();
        }
    }

public:
    void init() override
    {
        // Set up keybinding and overlay effect
        this->output->add_key(group_key, &group_key_callback);
        this->output->add_button(click_binding, &click_callback);
        this->output->render->add_effect(&overlay_effect, wf::OUTPUT_EFFECT_OVERLAY);

        // Connect to view signals
        wf::get_core().connect(&on_view_mapped);
        wf::get_core().connect(&on_view_unmapped);

        // Connect to geometry changes for existing views
        for (auto& view : wf::get_core().get_all_views())
        {
            if (wf::toplevel_cast(view))
            {
                view->connect(&on_view_geometry_changed);
                if (auto toplevel = wf::toplevel_cast(view))
                {
                    last_positions[view] = toplevel->get_geometry();
                }
            }
        }
    }

    void create_group_from_focused(wayfire_view anchor_view)
    {
        auto toplevel_anchor = wf::toplevel_cast(anchor_view);
        if (!toplevel_anchor) return;

        LOGD("GROUP: Creating new group with anchor ", anchor_view->get_title());

        auto new_group = std::make_shared<window_group_t>();
        auto anchor_geometry = toplevel_anchor->get_geometry();

        // Add the anchor view to the group (not hidden)
        group_member_t anchor_member;
        anchor_member.view = anchor_view;
        anchor_member.transformer = nullptr;
        anchor_member.is_hidden = false;
        new_group->members.emplace_back(std::move(anchor_member));
        view_to_group_map[anchor_view] = new_group;

        // Find overlapping windows and group them
        for (auto& view : wf::get_core().get_all_views())
        {
            auto toplevel_other = wf::toplevel_cast(view);
            if (!toplevel_other || view == anchor_view || !view->is_mapped())
            {
                continue;
            }

            auto view_geometry = toplevel_other->get_geometry();
            if (boxes_intersect(anchor_geometry, view_geometry))
            {
                group_member_t member;
                member.view = view;
                member.is_hidden = false;

                // Create transformer for movement
                member.transformer = std::make_shared<wf::scene::view_2d_transformer_t>(view);
                view->get_transformed_node()->add_transformer(member.transformer,
                    wf::TRANSFORMER_2D, group_transformer_name);

                // Move to anchor position
                int diff_x = anchor_geometry.x - view_geometry.x;
                int diff_y = anchor_geometry.y - view_geometry.y;
                member.transformer->translation_x = diff_x;
                member.transformer->translation_y = diff_y;

                // Hide the grouped window and create preview
                hide_window(member);
                create_preview_overlay(member, anchor_view);

                new_group->members.emplace_back(std::move(member));
                view_to_group_map[view] = new_group;

                LOGD("GROUP: Added ", view->get_title(), " to group.");
            }
        }

        update_preview_positions(new_group);
    }

    void ungroup_view(wayfire_view view)
    {
        if (!view_to_group_map.count(view)) return;

        LOGD("GROUP: Dissolving group for view ", view->get_title());

        auto group = view_to_group_map[view];

        for (auto& member : group->members)
        {
            if (member.transformer)
            {
                show_window(member);
                member.view->get_transformed_node()->rem_transformer(group_transformer_name);
                member.view->damage();
            }

            destroy_preview_overlay(member);
            view_to_group_map.erase(member.view);
        }

        // Schedule redraw to clear overlays
        this->output->render->schedule_redraw();
    }

    void update_group_transformers(wayfire_view moved_view)
    {
        auto group_it = view_to_group_map.find(moved_view);
        if (group_it == view_to_group_map.end()) return;

        auto group = group_it->second;
        auto toplevel_moved = wf::toplevel_cast(moved_view);
        if (!toplevel_moved) return;

        auto current_geometry = toplevel_moved->get_geometry();
        auto last_geometry = last_positions[moved_view];

        int delta_x = current_geometry.x - last_geometry.x;
        int delta_y = current_geometry.y - last_geometry.y;

        for (auto& member : group->members)
        {
            if (member.view == moved_view || !member.transformer) continue;

            member.transformer->translation_x += delta_x;
            member.transformer->translation_y += delta_y;
            member.view->damage();
        }

        last_positions[moved_view] = current_geometry;
        update_preview_positions(group);
    }

    // Signal handlers
    wf::signal::connection_t<wf::view_geometry_changed_signal> on_view_geometry_changed =
        [=] (wf::view_geometry_changed_signal *ev)
    {
        wayfire_view moved_view = ev->view;

        if (auto toplevel = wf::toplevel_cast(moved_view))
        {
            auto current_geometry = toplevel->get_geometry();

            if (view_to_group_map.count(moved_view))
            {
                if (last_positions.count(moved_view))
                {
                    auto last_geometry = last_positions[moved_view];
                    if (current_geometry.x != last_geometry.x || current_geometry.y != last_geometry.y)
                    {
                        update_group_transformers(moved_view);
                    }
                }
                else
                {
                    last_positions[moved_view] = current_geometry;
                }
            }
            else
            {
                last_positions[moved_view] = current_geometry;
            }
        }
    };

    wf::signal::connection_t<wf::view_mapped_signal> on_view_mapped =
        [=] (wf::view_mapped_signal *ev)
    {
        if (wf::toplevel_cast(ev->view))
        {
            ev->view->connect(&on_view_geometry_changed);
            if (auto toplevel = wf::toplevel_cast(ev->view))
            {
                last_positions[ev->view] = toplevel->get_geometry();
            }
        }
    };

    wf::signal::connection_t<wf::view_unmapped_signal> on_view_unmapped =
        [=] (wf::view_unmapped_signal *ev)
    {
        last_positions.erase(ev->view);

        if (view_to_group_map.count(ev->view))
        {
            auto group = view_to_group_map[ev->view];

            // Destroy preview overlay for this view
            for (auto& member : group->members)
            {
                if (member.view == ev->view)
                {
                    destroy_preview_overlay(member);
                    break;
                }
            }

            // Remove from group
            group->members.erase(
                std::remove_if(group->members.begin(), group->members.end(),
                    [=](const group_member_t& member) { return member.view == ev->view; }),
                group->members.end()
            );

            view_to_group_map.erase(ev->view);

            // Dissolve group if too few members
            if (group->members.size() <= 1)
            {
                if (!group->members.empty())
                {
                    auto& last_member = group->members.front();
                    if (last_member.transformer)
                    {
                        show_window(last_member);
                        last_member.view->get_transformed_node()->rem_transformer(group_transformer_name);
                        last_member.view->damage();
                    }
                    destroy_preview_overlay(last_member);
                    view_to_group_map.erase(last_member.view);
                }
            }
            else
            {
                update_preview_positions(group);
            }
        }
    };

    void fini() override
    {
        // Clean up
        this->output->rem_binding(&group_key_callback);
        this->output->rem_binding(&click_callback);
        this->output->render->rem_effect(&overlay_effect);

        // Remove all transformers and preview overlays
        for (const auto& [view, group] : view_to_group_map)
        {
            for (auto& member : group->members)
            {
                if (member.transformer)
                {
                    show_window(member);
                    member.view->get_transformed_node()->rem_transformer(group_transformer_name);
                }
                destroy_preview_overlay(member);
            }
        }

        // Disconnect signals
        on_view_mapped.disconnect();
        on_view_unmapped.disconnect();
        on_view_geometry_changed.disconnect();

        // Clear tracking
        last_positions.clear();
        view_to_group_map.clear();
    }
};

class wayfire_group_global : public wf::plugin_interface_t,
    public wf::per_output_tracker_mixin_t<wayfire_group>
{
  public:
    void init() override
    {
        this->init_output_tracking();
    }

    void fini() override
    {
        this->fini_output_tracking();
    }
};

DECLARE_WAYFIRE_PLUGIN(wayfire_group_global);