TreeWidget.h

Go to the documentation of this file.

#pragma once
 
#include "Core/BehaviorTree.h"
#include "Widgets/NodeWidget.h"
#include <memory>
#include <unordered_map>
#include <vector>
#include <wx/scrolwin.h>
#include <wx/wx.h>
 
namespace EmberForge {

class TreeWidget : public wxScrolledWindow {
  public:
    enum class LayoutAlgorithm {
        Vertical,   // Top-down tree layout
        Horizontal, // Left-right tree layout
        Radial,     // Circular tree layout
        Compact,    // Space-efficient layout
        Manual      // User-controlled positions
    };

    enum class InteractionMode {
        View,    // Read-only viewing
        Select,  // Selection mode
        Edit,    // Editing mode (add/remove/move nodes)
        Connect, // Connection editing mode
        Debug    // Debug/execution visualization
    };

    TreeWidget(wxWindow *parent, std::shared_ptr<EmberCore::BehaviorTree> tree = nullptr);

    virtual ~TreeWidget();
 
    // Tree management
    void SetBehaviorTree(std::shared_ptr<EmberCore::BehaviorTree> tree);
    std::shared_ptr<EmberCore::BehaviorTree> GetBehaviorTree() const { return behavior_tree_; }
    bool HasTree() const { return behavior_tree_ != nullptr; }
 
    // Layout management
    LayoutAlgorithm GetLayoutAlgorithm() const { return layout_algorithm_; }
    void SetLayoutAlgorithm(LayoutAlgorithm algorithm);
    void RefreshLayout();
    void AutoFitLayout();
 
    // Interaction mode
    InteractionMode GetInteractionMode() const { return interaction_mode_; }
    void SetInteractionMode(InteractionMode mode);
 
    // EmberCore::Node management
    EmberForge::NodeWidget *GetNodeWidget(EmberCore::Node *node) const;
    std::vector<EmberForge::NodeWidget *> GetAllNodeWidgets() const;
    std::vector<EmberForge::NodeWidget *> GetSelectedNodeWidgets() const;
 
    // Selection
    void SelectNode(EmberCore::Node *node, bool add_to_selection = false);
    void DeselectNode(EmberCore::Node *node);
    void ClearSelection();
    void SelectAll();
    EmberCore::Node *GetSelectedNode() const; // Returns first selected node
    std::vector<EmberCore::Node *> GetSelectedNodes() const;
 
    // Navigation
    void ZoomIn();
    void ZoomOut();
    void ZoomToFit();
    void ZoomToSelection();
    void SetZoomLevel(double zoom);
    double GetZoomLevel() const { return zoom_level_; }
 
    void CenterOnNode(EmberCore::Node *node);
    void CenterOnSelection();
    void CenterView();
 
    // Tree editing
    bool CanAddNode(EmberCore::Node *parent) const;
    bool AddNode(std::unique_ptr<EmberCore::Node> node, EmberCore::Node *parent);
    bool RemoveSelectedNodes();
    bool DuplicateSelectedNodes();
    bool CanMoveNode(EmberCore::Node *node, EmberCore::Node *new_parent) const;
    bool MoveNode(EmberCore::Node *node, EmberCore::Node *new_parent);
 
    // Visual properties
    void SetNodeRenderStyle(EmberForge::NodeWidget::RenderStyle style);
    void SetConnectionStyle(int line_width, const wxColour &color);
    void SetBackgroundColor(const wxColour &color);
    void SetGridVisible(bool visible);
    void SetConnectionsVisible(bool visible);
 
    // Animation
    void AnimateToNode(EmberCore::Node *node);
    void HighlightExecutionPath(const std::vector<EmberCore::Node *> &path);
    void ClearHighlights();
 
    // Import/Export
    bool SaveLayout(const wxString &filename) const;
    bool LoadLayout(const wxString &filename);
    wxBitmap ExportAsImage(const wxSize &size = wxDefaultSize) const;
 
    // Preferences
    void LoadPreferences();
 
    // Event callbacks
    using NodeSelectionCallback = std::function<void(const std::vector<EmberCore::Node *> &)>;
    using NodeEditCallback = std::function<void(EmberCore::Node *, const wxString &property, const wxString &value)>;
    using TreeChangeCallback = std::function<void(const wxString &change_type)>;
 
    void SetNodeSelectionCallback(NodeSelectionCallback callback) { selection_callback_ = callback; }
    void SetNodeEditCallback(NodeEditCallback callback) { edit_callback_ = callback; }
    void SetTreeChangeCallback(TreeChangeCallback callback) { tree_change_callback_ = callback; }
 
  protected:
    // Rendering
    virtual void OnPaint(wxPaintEvent &event);
    virtual void OnSize(wxSizeEvent &event);
    virtual void OnEraseBackground(wxEraseEvent &event);
 
    // Mouse events
    virtual void OnMouseDown(wxMouseEvent &event);
    virtual void OnMouseUp(wxMouseEvent &event);
    virtual void OnMouseMove(wxMouseEvent &event);
    virtual void OnMouseWheel(wxMouseEvent &event);
    virtual void OnRightClick(wxContextMenuEvent &event);
 
    // Keyboard events
    virtual void OnKeyDown(wxKeyEvent &event);
    virtual void OnChar(wxKeyEvent &event);
 
    // Layout methods
    void CalculateVerticalLayout();
    void CalculateHorizontalLayout();
    void CalculateRadialLayout();
    void CalculateCompactLayout();
    void UpdateScrollbars();
    void LayoutNodeVertical(EmberCore::Node *node, int x, int y, int level);
 
    // Drawing methods
    void DrawBackground(wxDC &dc);
    void DrawGrid(wxDC &dc);
    void DrawConnections(wxDC &dc);
    void DrawConnection(wxDC &dc, EmberForge::NodeWidget *parent, EmberForge::NodeWidget *child);
    void DrawSelectionRect(wxDC &dc);
 
    // Node widget management
    void CreateNodeWidgets();
    void UpdateNodeWidgets();
    void DestroyNodeWidgets();
    std::unique_ptr<EmberForge::NodeWidget> CreateNodeWidget(EmberCore::Node *node);
 
    // Hit testing
    EmberForge::NodeWidget *HitTestNode(const wxPoint &point) const;
    EmberCore::Node *HitTestConnection(const wxPoint &point) const;
 
    // Selection helpers
    void UpdateSelectionVisuals();
    void NotifySelectionChange();
 
    // Layout helpers
    wxSize CalculateTreeSize() const;
    wxPoint CalculateNodePosition(EmberCore::Node *node) const;
    void PositionNodeWidgets();
 
    // Coordinate conversion
    wxPoint TreeToScreen(const wxPoint &tree_point) const;
    wxPoint ScreenToTree(const wxPoint &screen_point) const;
    wxRect TreeToScreen(const wxRect &tree_rect) const;
    wxRect ScreenToTree(const wxRect &screen_rect) const;
 
  private:
    // Core data
    std::shared_ptr<EmberCore::BehaviorTree> behavior_tree_;
    LayoutAlgorithm layout_algorithm_;
    InteractionMode interaction_mode_;
 
    // Node widgets
    std::unordered_map<EmberCore::Node *, std::unique_ptr<EmberForge::NodeWidget>> node_widgets_;
    std::vector<EmberCore::Node *> selected_nodes_;
 
    // Visual state
    double zoom_level_;
    wxPoint view_offset_;
    wxSize tree_size_;
    EmberForge::NodeWidget::RenderStyle node_render_style_;
 
    // Zoom/Pan preferences
    double zoom_step_;
    double mouse_wheel_sensitivity_;
    double pan_sensitivity_;
    bool zoom_follows_cursor_;
 
    // Appearance
    wxColour background_color_;
    wxColour grid_color_;
    wxColour connection_color_;
    int connection_width_;
    bool grid_visible_;
    bool connections_visible_;
 
    // Interaction state
    bool is_dragging_;
    bool is_selecting_;
    wxPoint drag_start_pos_;
    wxPoint selection_start_pos_;
    wxRect selection_rect_;
 
    // Layout parameters
    int node_spacing_x_;
    int node_spacing_y_;
    int level_spacing_;
    wxPoint tree_root_position_;
    std::unordered_map<EmberCore::Node *, wxPoint> node_positions_;
 
    // Animation
    wxTimer *animation_timer_;
    std::vector<EmberCore::Node *> highlighted_path_;
 
    // Event callbacks
    NodeSelectionCallback selection_callback_;
    NodeEditCallback edit_callback_;
    TreeChangeCallback tree_change_callback_;
 
    // Helper methods
    void InitializeWidget();
    void SetupEventHandlers();
    void UpdateViewSize();
    void ShowContextMenu(const wxPoint &position);
 
    // Node widget event handlers
    void OnNodeClicked(EmberForge::NodeWidget *widget, EmberCore::Node *node);
    void OnNodeDoubleClicked(EmberForge::NodeWidget *widget, EmberCore::Node *node);
    void OnNodeDragged(EmberForge::NodeWidget *widget, const wxPoint &delta);
    void OnNodeEdited(EmberForge::NodeWidget *widget, const wxString &new_name);
 
    // Constants
    static constexpr double MIN_ZOOM = 0.1;
    static constexpr double MAX_ZOOM = 5.0;
    static const int DEFAULT_NODE_SPACING_X = 120;
    static const int DEFAULT_NODE_SPACING_Y = 80;
    static const int DEFAULT_LEVEL_SPACING = 100;
    static const int GRID_SIZE = 20;
 
    DECLARE_EVENT_TABLE()
};
 
} // namespace EmberForge