warp-pipe/gui/GraphEditorWidget.cpp

#include "GraphEditorWidget.h"
#include "WarpGraphModel.h"

#include <QtNodes/BasicGraphicsScene>
#include <QtNodes/ConnectionStyle>
#include <QtNodes/GraphicsView>

#include <QDateTime>
#include <QDir>
#include <QInputDialog>
#include <QMenu>
#include <QMessageBox>
#include <QPixmap>
#include <QStandardPaths>
#include <QTimer>
#include <QVBoxLayout>

GraphEditorWidget::GraphEditorWidget(warppipe::Client *client,
                                     QWidget *parent)
    : QWidget(parent), m_client(client) {
  m_layoutPath =
      QStandardPaths::writableLocation(QStandardPaths::AppConfigLocation) +
      QStringLiteral("/layout.json");

  m_model = new WarpGraphModel(client, this);
  bool hasLayout = m_model->loadLayout(m_layoutPath);

  m_scene = new QtNodes::BasicGraphicsScene(*m_model, this);

  QtNodes::ConnectionStyle::setConnectionStyle(
      R"({"ConnectionStyle": {
        "ConstructionColor": "#b4b4c8",
        "NormalColor": "#c8c8dc",
        "SelectedColor": "#ffa500",
        "SelectedHaloColor": "#ffa50040",
        "HoveredColor": "#f0c878",
        "LineWidth": 2.4,
        "ConstructionLineWidth": 1.8,
        "PointDiameter": 10.0,
        "UseDataDefinedColors": false
      }})");

  m_view = new QtNodes::GraphicsView(m_scene);

  auto *layout = new QVBoxLayout(this);
  layout->setContentsMargins(0, 0, 0, 0);
  layout->addWidget(m_view);

  m_view->setContextMenuPolicy(Qt::CustomContextMenu);
  connect(m_view, &QWidget::customContextMenuRequested, this,
          &GraphEditorWidget::onContextMenuRequested);

  connect(m_model, &WarpGraphModel::nodePositionUpdated, this,
          &GraphEditorWidget::scheduleSaveLayout);

  m_saveTimer = new QTimer(this);
  m_saveTimer->setSingleShot(true);
  m_saveTimer->setInterval(1000);
  connect(m_saveTimer, &QTimer::timeout, this, [this]() {
    m_model->saveLayout(m_layoutPath);
  });

  m_model->refreshFromClient();
  if (!hasLayout) {
    m_model->autoArrange();
  }

  if (m_model->allNodeIds().size() > 0) {
    m_graphReady = true;
    Q_EMIT graphReady();
  }

  m_refreshTimer = new QTimer(this);
  connect(m_refreshTimer, &QTimer::timeout, this,
          &GraphEditorWidget::onRefreshTimer);
  m_refreshTimer->start(500);
}

void GraphEditorWidget::onRefreshTimer() {
  m_model->refreshFromClient();

  if (!m_graphReady && m_model->allNodeIds().size() > 0) {
    m_graphReady = true;
    Q_EMIT graphReady();
    captureDebugScreenshot("initial_load");
  }
}

void GraphEditorWidget::scheduleSaveLayout() {
  if (!m_saveTimer->isActive()) {
    m_saveTimer->start();
  }
}

int GraphEditorWidget::nodeCount() const {
  return static_cast<int>(m_model->allNodeIds().size());
}

int GraphEditorWidget::linkCount() const {
  int count = 0;
  for (auto nodeId : m_model->allNodeIds()) {
    count += static_cast<int>(
        m_model->allConnectionIds(nodeId).size());
  }
  return count / 2;
}

void GraphEditorWidget::setDebugScreenshotDir(const QString &dir) {
  m_debugScreenshotDir = dir;
  QDir d(dir);
  if (!d.exists()) {
    d.mkpath(".");
  }

  connect(m_model, &QtNodes::AbstractGraphModel::nodeCreated, this, [this]() {
    captureDebugScreenshot("node_added");
  });
  connect(m_model, &QtNodes::AbstractGraphModel::nodeDeleted, this, [this]() {
    captureDebugScreenshot("node_removed");
  });
  connect(m_model, &QtNodes::AbstractGraphModel::connectionCreated, this,
          [this]() { captureDebugScreenshot("connection_added"); });
  connect(m_model, &QtNodes::AbstractGraphModel::connectionDeleted, this,
          [this]() { captureDebugScreenshot("connection_removed"); });
  connect(m_model, &QtNodes::AbstractGraphModel::nodeUpdated, this, [this]() {
    captureDebugScreenshot("node_updated");
  });

  if (m_graphReady) {
    QTimer::singleShot(200, this, [this]() {
      captureDebugScreenshot("initial_load");
    });
  }
}

void GraphEditorWidget::captureDebugScreenshot(const QString &event) {
  if (m_debugScreenshotDir.isEmpty()) {
    return;
  }

  QWidget *win = window();
  if (!win) {
    return;
  }

  QPixmap pixmap = win->grab();
  if (pixmap.isNull()) {
    return;
  }

  QString timestamp =
      QDateTime::currentDateTime().toString("yyyyMMdd_HHmmss");
  QString filename = QString("warppipe_%1_%2.png").arg(timestamp, event);
  pixmap.save(m_debugScreenshotDir + "/" + filename);
}

void GraphEditorWidget::onContextMenuRequested(const QPoint &pos) {
  QPointF scenePos = m_view->mapToScene(pos);

  uint32_t hitPwNodeId = 0;
  for (auto nodeId : m_model->allNodeIds()) {
    const WarpNodeData *data = m_model->warpNodeData(nodeId);
    if (!data) {
      continue;
    }
    QPointF nodePos =
        m_model->nodeData(nodeId, QtNodes::NodeRole::Position).toPointF();
    QSize nodeSize =
        m_model->nodeData(nodeId, QtNodes::NodeRole::Size).toSize();
    QRectF nodeRect(nodePos, QSizeF(nodeSize));
    if (nodeRect.contains(scenePos)) {
      hitPwNodeId = data->info.id.value;
      break;
    }
  }

  QPoint screenPos = m_view->mapToGlobal(pos);
  if (hitPwNodeId != 0) {
    showNodeContextMenu(screenPos, hitPwNodeId);
  } else {
    showCanvasContextMenu(screenPos, scenePos);
  }
}

void GraphEditorWidget::showCanvasContextMenu(const QPoint &screenPos,
                                               const QPointF &scenePos) {
  QMenu menu;
  QAction *createSink = menu.addAction(QStringLiteral("Create Virtual Sink"));
  QAction *createSource =
      menu.addAction(QStringLiteral("Create Virtual Source"));
  menu.addSeparator();
  QAction *autoArrange = menu.addAction(QStringLiteral("Auto-Arrange"));

  QAction *chosen = menu.exec(screenPos);
  if (chosen == createSink) {
    createVirtualNode(true, scenePos);
  } else if (chosen == createSource) {
    createVirtualNode(false, scenePos);
  } else if (chosen == autoArrange) {
    m_model->autoArrange();
  }
}

void GraphEditorWidget::showNodeContextMenu(const QPoint &screenPos,
                                            uint32_t pwNodeId) {
  QtNodes::NodeId qtId = m_model->qtNodeIdForPw(pwNodeId);
  const WarpNodeData *data = m_model->warpNodeData(qtId);
  if (!data) {
    return;
  }

  WarpNodeType type = WarpGraphModel::classifyNode(data->info);
  bool isVirtual =
      type == WarpNodeType::kVirtualSink || type == WarpNodeType::kVirtualSource;

  if (!isVirtual) {
    return;
  }

  QMenu menu;
  QAction *deleteAction = menu.addAction(QStringLiteral("Delete Node"));

  QAction *chosen = menu.exec(screenPos);
  if (chosen == deleteAction && m_client) {
    m_client->RemoveNode(warppipe::NodeId{pwNodeId});
    m_model->refreshFromClient();
  }
}

void GraphEditorWidget::createVirtualNode(bool isSink,
                                          const QPointF &scenePos) {
  QString label = isSink ? QStringLiteral("Create Virtual Sink")
                         : QStringLiteral("Create Virtual Source");
  bool ok = false;
  QString name = QInputDialog::getText(this, label,
                                       QStringLiteral("Node name:"),
                                       QLineEdit::Normal, QString(), &ok);
  if (!ok || name.trimmed().isEmpty()) {
    return;
  }

  std::string nodeName = name.trimmed().toStdString();
  m_model->setPendingPosition(nodeName, scenePos);

  warppipe::Status status;
  if (isSink) {
    auto result = m_client->CreateVirtualSink(nodeName);
    status = result.status;
  } else {
    auto result = m_client->CreateVirtualSource(nodeName);
    status = result.status;
  }

  if (!status.ok()) {
    QMessageBox::warning(this, QStringLiteral("Error"),
                         QString::fromStdString(status.message));
    return;
  }

  m_model->refreshFromClient();
}