warp-pipe/gui/WarpGraphModel.h

#pragma once

#include <warppipe/warppipe.hpp>

#include <QtNodes/AbstractGraphModel>
#include <QtNodes/ConnectionIdUtils>

#include <QHash>
#include <QPointF>
#include <QSize>
#include <QString>

#include <unordered_map>
#include <unordered_set>

enum class WarpNodeType : uint8_t {
  kUnknown = 0,
  kHardwareSink,
  kHardwareSource,
  kVirtualSink,
  kVirtualSource,
  kApplication,
  kVideoSource,
  kVideoSink,
};

inline bool nodeHasVolume(WarpNodeType type) {
  switch (type) {
  case WarpNodeType::kVideoSource:
  case WarpNodeType::kVideoSink:
  case WarpNodeType::kUnknown:
    return false;
  default:
    return true;
  }
}

struct WarpNodeData {
  warppipe::NodeInfo info;
  std::vector<warppipe::PortInfo> inputPorts;
  std::vector<warppipe::PortInfo> outputPorts;
};

class WarpGraphModel : public QtNodes::AbstractGraphModel {
  Q_OBJECT

public:
  explicit WarpGraphModel(warppipe::Client *client, QObject *parent = nullptr);

  QtNodes::NodeId newNodeId() override;
  std::unordered_set<QtNodes::NodeId> allNodeIds() const override;
  std::unordered_set<QtNodes::ConnectionId> allConnectionIds(
      QtNodes::NodeId nodeId) const override;
  std::unordered_set<QtNodes::ConnectionId> connections(
      QtNodes::NodeId nodeId, QtNodes::PortType portType,
      QtNodes::PortIndex portIndex) const override;
  bool connectionExists(QtNodes::ConnectionId const connectionId) const override;
  QtNodes::NodeId addNode(QString const nodeType = QString()) override;
  bool connectionPossible(QtNodes::ConnectionId const connectionId) const override;
  void addConnection(QtNodes::ConnectionId const connectionId) override;
  bool nodeExists(QtNodes::NodeId const nodeId) const override;
  QVariant nodeData(QtNodes::NodeId nodeId, QtNodes::NodeRole role) const override;
  bool setNodeData(QtNodes::NodeId nodeId, QtNodes::NodeRole role,
                   QVariant value) override;
  QVariant portData(QtNodes::NodeId nodeId, QtNodes::PortType portType,
                    QtNodes::PortIndex portIndex,
                    QtNodes::PortRole role) const override;
  bool setPortData(QtNodes::NodeId nodeId, QtNodes::PortType portType,
                   QtNodes::PortIndex portIndex, QVariant const &value,
                   QtNodes::PortRole role = QtNodes::PortRole::Data) override;
  bool deleteConnection(QtNodes::ConnectionId const connectionId) override;
  bool deleteNode(QtNodes::NodeId const nodeId) override;
  QJsonObject saveNode(QtNodes::NodeId const) const override;
  void loadNode(QJsonObject const &) override;

  void refreshFromClient();
  const WarpNodeData *warpNodeData(QtNodes::NodeId nodeId) const;
  QtNodes::NodeId qtNodeIdForPw(uint32_t pwNodeId) const;
  bool isGhost(QtNodes::NodeId nodeId) const;
  void setPendingPosition(const std::string &nodeName, QPointF pos);
  static WarpNodeType classifyNode(const warppipe::NodeInfo &info);

  uint32_t findPwNodeIdByName(const std::string &name) const;

  struct NodeVolumeState {
    float volume = 1.0f;
    bool mute = false;
  };

  void setNodeVolumeState(QtNodes::NodeId nodeId, const NodeVolumeState &state);
  NodeVolumeState nodeVolumeState(QtNodes::NodeId nodeId) const;

Q_SIGNALS:
  void nodeVolumeChanged(QtNodes::NodeId nodeId, NodeVolumeState previous,
                         NodeVolumeState current);

public:
  struct ViewState {
    double scale;
    double centerX;
    double centerY;
    double zoomSensitivity;
    double zoomMin;
    double zoomMax;
    int splitterGraph;
    int splitterSidebar;
    int connectionStyle;
    bool valid;
  };

  void saveLayout(const QString &path) const;
  void saveLayout(const QString &path, const ViewState &viewState) const;
  bool loadLayout(const QString &path);
  ViewState savedViewState() const;
  void clearSavedPositions();
  void autoArrange();

private:
  static QString captionForNode(const warppipe::NodeInfo &info);
  static QVariant styleForNode(WarpNodeType type, bool ghost);
  QPointF nextPosition(const WarpNodeData &data);
  QPointF findNonOverlappingPosition(QPointF candidate, const WarpNodeData &data) const;
  static QSize estimateNodeSize(const WarpNodeData &data);

  warppipe::Client *m_client = nullptr;
  QtNodes::NodeId m_nextNodeId = 1;

  std::unordered_map<QtNodes::NodeId, WarpNodeData> m_nodes;
  std::unordered_map<uint32_t, QtNodes::NodeId> m_pwToQt;

  std::unordered_set<QtNodes::ConnectionId> m_connections;
  std::unordered_map<uint32_t, QtNodes::ConnectionId> m_linkIdToConn;

  std::unordered_map<QtNodes::NodeId, QPointF> m_positions;
  std::unordered_map<QtNodes::NodeId, QSize> m_sizes;
  std::unordered_set<QtNodes::NodeId> m_ghostNodes;
  std::unordered_set<QtNodes::ConnectionId> m_ghostConnections;

  static constexpr double kHorizontalGap = 40.0;
  static constexpr double kVerticalGap = 30.0;
  static constexpr double kMaxRowWidth = 1400.0;
  double m_nextX = 0.0;
  double m_nextY = 0.0;
  double m_rowMaxHeight = 0.0;

  bool m_refreshing = false;

  struct PendingGhostConnection {
    std::string outNodeName;
    std::string outPortName;
    std::string inNodeName;
    std::string inPortName;
  };

  std::unordered_map<std::string, QPointF> m_pendingPositions;
  std::unordered_map<std::string, QPointF> m_savedPositions;
  std::vector<PendingGhostConnection> m_pendingGhostConnections;
  ViewState m_savedViewState{};

  std::unordered_map<QtNodes::NodeId, NodeVolumeState> m_volumeStates;
  std::unordered_map<QtNodes::NodeId, QWidget *> m_volumeWidgets;
  mutable std::unordered_map<QtNodes::NodeId, QVariant> m_styleCache;
};