#include "WarpGraphModel.h"

#include <QColor>
#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QJsonArray>
#include <QJsonDocument>
#include <QJsonObject>
#include <QVariant>

#include <QtNodes/NodeStyle>
#include <QtNodes/StyleCollection>

#include <algorithm>
#include <cmath>

WarpGraphModel::WarpGraphModel(warppipe::Client *client, QObject *parent)
    : QtNodes::AbstractGraphModel(), m_client(client) {
  if (parent) {
    setParent(parent);
  }
}

QtNodes::NodeId WarpGraphModel::newNodeId() { return m_nextNodeId++; }

std::unordered_set<QtNodes::NodeId> WarpGraphModel::allNodeIds() const {
  std::unordered_set<QtNodes::NodeId> ids;
  ids.reserve(m_nodes.size());
  for (const auto &entry : m_nodes) {
    ids.insert(entry.first);
  }
  return ids;
}

std::unordered_set<QtNodes::ConnectionId>
WarpGraphModel::allConnectionIds(QtNodes::NodeId const nodeId) const {
  std::unordered_set<QtNodes::ConnectionId> result;
  for (const auto &conn : m_connections) {
    if (conn.outNodeId == nodeId || conn.inNodeId == nodeId) {
      result.insert(conn);
    }
  }
  return result;
}

std::unordered_set<QtNodes::ConnectionId>
WarpGraphModel::connections(QtNodes::NodeId nodeId,
                            QtNodes::PortType portType,
                            QtNodes::PortIndex portIndex) const {
  std::unordered_set<QtNodes::ConnectionId> result;
  for (const auto &conn : m_connections) {
    if (portType == QtNodes::PortType::Out) {
      if (conn.outNodeId == nodeId && conn.outPortIndex == portIndex) {
        result.insert(conn);
      }
    } else if (portType == QtNodes::PortType::In) {
      if (conn.inNodeId == nodeId && conn.inPortIndex == portIndex) {
        result.insert(conn);
      }
    }
  }
  return result;
}

bool WarpGraphModel::connectionExists(
    QtNodes::ConnectionId const connectionId) const {
  return m_connections.find(connectionId) != m_connections.end();
}

QtNodes::NodeId WarpGraphModel::addNode(QString const) {
  return newNodeId();
}

bool WarpGraphModel::connectionPossible(
    QtNodes::ConnectionId const connectionId) const {
  if (!nodeExists(connectionId.outNodeId) ||
      !nodeExists(connectionId.inNodeId)) {
    return false;
  }
  if (connectionExists(connectionId)) {
    return false;
  }

  auto outIt = m_nodes.find(connectionId.outNodeId);
  auto inIt = m_nodes.find(connectionId.inNodeId);
  if (outIt == m_nodes.end() || inIt == m_nodes.end()) {
    return false;
  }

  auto outIdx = static_cast<size_t>(connectionId.outPortIndex);
  auto inIdx = static_cast<size_t>(connectionId.inPortIndex);
  if (outIdx >= outIt->second.outputPorts.size()) {
    return false;
  }
  if (inIdx >= inIt->second.inputPorts.size()) {
    return false;
  }

  return true;
}

void WarpGraphModel::addConnection(
    QtNodes::ConnectionId const connectionId) {
  if (!connectionPossible(connectionId)) {
    return;
  }

  if (m_client) {
    auto outIt = m_nodes.find(connectionId.outNodeId);
    auto inIt = m_nodes.find(connectionId.inNodeId);
    if (outIt == m_nodes.end() || inIt == m_nodes.end()) {
      return;
    }

    auto outIdx = static_cast<size_t>(connectionId.outPortIndex);
    auto inIdx = static_cast<size_t>(connectionId.inPortIndex);
    if (outIdx >= outIt->second.outputPorts.size() ||
        inIdx >= inIt->second.inputPorts.size()) {
      return;
    }

    warppipe::PortId outPortId = outIt->second.outputPorts[outIdx].id;
    warppipe::PortId inPortId = inIt->second.inputPorts[inIdx].id;

    auto result = m_client->CreateLink(outPortId, inPortId, warppipe::LinkOptions{});
    if (!result.ok()) {
      return;
    }

    m_linkIdToConn.emplace(result.value.id.value, connectionId);
  }

  m_connections.insert(connectionId);
  Q_EMIT connectionCreated(connectionId);
}

bool WarpGraphModel::nodeExists(QtNodes::NodeId const nodeId) const {
  return m_nodes.find(nodeId) != m_nodes.end();
}

QVariant WarpGraphModel::nodeData(QtNodes::NodeId nodeId,
                                  QtNodes::NodeRole role) const {
  auto it = m_nodes.find(nodeId);
  if (it == m_nodes.end()) {
    return QVariant();
  }

  const auto &data = it->second;

  switch (role) {
  case QtNodes::NodeRole::Caption:
    return captionForNode(data.info);
  case QtNodes::NodeRole::CaptionVisible:
    return true;
  case QtNodes::NodeRole::Position: {
    auto posIt = m_positions.find(nodeId);
    if (posIt != m_positions.end()) {
      return posIt->second;
    }
    return QPointF(0, 0);
  }
  case QtNodes::NodeRole::Size: {
    auto sizeIt = m_sizes.find(nodeId);
    if (sizeIt != m_sizes.end()) {
      return sizeIt->second;
    }
    return estimateNodeSize(data);
  }
  case QtNodes::NodeRole::InPortCount:
    return static_cast<unsigned int>(data.inputPorts.size());
  case QtNodes::NodeRole::OutPortCount:
    return static_cast<unsigned int>(data.outputPorts.size());
  case QtNodes::NodeRole::Type:
    return QString("PipeWire");
  case QtNodes::NodeRole::Style: {
    bool ghost = m_ghostNodes.find(nodeId) != m_ghostNodes.end();
    WarpNodeType type = classifyNode(data.info);
    return styleForNode(type, ghost);
  }
  default:
    return QVariant();
  }
}

bool WarpGraphModel::setNodeData(QtNodes::NodeId nodeId,
                                 QtNodes::NodeRole role, QVariant value) {
  if (!nodeExists(nodeId)) {
    return false;
  }

  if (role == QtNodes::NodeRole::Position) {
    m_positions[nodeId] = value.toPointF();
    Q_EMIT nodePositionUpdated(nodeId);
    return true;
  }

  if (role == QtNodes::NodeRole::Size) {
    m_sizes[nodeId] = value.toSize();
    return true;
  }

  return false;
}

QVariant WarpGraphModel::portData(QtNodes::NodeId nodeId,
                                  QtNodes::PortType portType,
                                  QtNodes::PortIndex portIndex,
                                  QtNodes::PortRole role) const {
  auto it = m_nodes.find(nodeId);
  if (it == m_nodes.end()) {
    return QVariant();
  }

  const auto &data = it->second;

  if (role == QtNodes::PortRole::DataType) {
    return QString("audio");
  }

  if (role == QtNodes::PortRole::CaptionVisible) {
    return true;
  }

  if (role == QtNodes::PortRole::Caption) {
    if (portType == QtNodes::PortType::In) {
      auto idx = static_cast<size_t>(portIndex);
      if (idx < data.inputPorts.size()) {
        return QString::fromStdString(data.inputPorts[idx].name);
      }
    } else if (portType == QtNodes::PortType::Out) {
      auto idx = static_cast<size_t>(portIndex);
      if (idx < data.outputPorts.size()) {
        return QString::fromStdString(data.outputPorts[idx].name);
      }
    }
  }

  if (role == QtNodes::PortRole::ConnectionPolicyRole) {
    return QVariant::fromValue(QtNodes::ConnectionPolicy::Many);
  }

  return QVariant();
}

bool WarpGraphModel::setPortData(QtNodes::NodeId, QtNodes::PortType,
                                 QtNodes::PortIndex, QVariant const &,
                                 QtNodes::PortRole) {
  return false;
}

bool WarpGraphModel::deleteConnection(
    QtNodes::ConnectionId const connectionId) {
  auto it = m_connections.find(connectionId);
  if (it == m_connections.end()) {
    return false;
  }

  if (m_client && !m_refreshing) {
    for (auto linkIt = m_linkIdToConn.begin(); linkIt != m_linkIdToConn.end();
         ++linkIt) {
      if (linkIt->second == connectionId) {
        m_client->RemoveLink(warppipe::LinkId{linkIt->first});
        m_linkIdToConn.erase(linkIt);
        break;
      }
    }
  }

  m_connections.erase(it);
  Q_EMIT connectionDeleted(connectionId);
  return true;
}

bool WarpGraphModel::deleteNode(QtNodes::NodeId const nodeId) {
  if (!nodeExists(nodeId)) {
    return false;
  }

  std::vector<QtNodes::ConnectionId> toRemove;
  for (const auto &conn : m_connections) {
    if (conn.outNodeId == nodeId || conn.inNodeId == nodeId) {
      toRemove.push_back(conn);
    }
  }
  for (const auto &conn : toRemove) {
    deleteConnection(conn);
  }

  m_nodes.erase(nodeId);
  m_positions.erase(nodeId);
  m_sizes.erase(nodeId);
  Q_EMIT nodeDeleted(nodeId);
  return true;
}

QJsonObject WarpGraphModel::saveNode(QtNodes::NodeId const nodeId) const {
  QJsonObject obj;
  obj["id"] = static_cast<qint64>(nodeId);
  QPointF pos = nodeData(nodeId, QtNodes::NodeRole::Position).toPointF();
  QJsonObject posObj;
  posObj["x"] = pos.x();
  posObj["y"] = pos.y();
  obj["position"] = posObj;
  return obj;
}

void WarpGraphModel::loadNode(QJsonObject const &) {}

void WarpGraphModel::refreshFromClient() {
  if (!m_client) {
    return;
  }

  m_refreshing = true;
  auto nodesResult = m_client->ListNodes();
  if (!nodesResult.ok()) {
    return;
  }

  std::unordered_set<uint32_t> seenPwIds;

  for (const auto &nodeInfo : nodesResult.value) {
    seenPwIds.insert(nodeInfo.id.value);

    WarpNodeType nodeType = classifyNode(nodeInfo);
    bool isStream = nodeType == WarpNodeType::kApplication;
    if (isStream && nodeInfo.name.empty() &&
        nodeInfo.application_name.empty()) {
      continue;
    }

    auto existing = m_pwToQt.find(nodeInfo.id.value);
    if (existing != m_pwToQt.end()) {
      QtNodes::NodeId qtId = existing->second;
      auto &data = m_nodes[qtId];
      data.info = nodeInfo;

      bool portsMissing =
          data.inputPorts.empty() && data.outputPorts.empty();
      if (portsMissing) {
        auto portsResult = m_client->ListPorts(nodeInfo.id);
        if (portsResult.ok() && !portsResult.value.empty()) {
          for (const auto &port : portsResult.value) {
            if (port.is_input) {
              data.inputPorts.push_back(port);
            } else {
              data.outputPorts.push_back(port);
            }
          }
          std::sort(data.inputPorts.begin(), data.inputPorts.end(),
                    [](const auto &a, const auto &b) {
                      return a.name < b.name;
                    });
          std::sort(data.outputPorts.begin(), data.outputPorts.end(),
                    [](const auto &a, const auto &b) {
                      return a.name < b.name;
                    });
          Q_EMIT nodeUpdated(qtId);
        }
      }

      if (m_ghostNodes.erase(qtId)) {
        std::erase_if(m_ghostConnections, [&](const auto &gc) {
          if (gc.outNodeId != qtId && gc.inNodeId != qtId)
            return false;
          m_connections.erase(gc);
          Q_EMIT connectionDeleted(gc);
          return true;
        });
        Q_EMIT nodeUpdated(qtId);
      }
      continue;
    }

    QtNodes::NodeId ghostMatch = 0;
    std::string nodeName = nodeInfo.name;
    for (const auto &ghostId : m_ghostNodes) {
      auto ghostIt = m_nodes.find(ghostId);
      if (ghostIt != m_nodes.end() &&
          ghostIt->second.info.name == nodeName) {
        ghostMatch = ghostId;
        break;
      }
    }

    if (ghostMatch != 0) {
      m_ghostNodes.erase(ghostMatch);
      std::erase_if(m_ghostConnections, [&](const auto &gc) {
        if (gc.outNodeId != ghostMatch && gc.inNodeId != ghostMatch)
          return false;
        m_connections.erase(gc);
        Q_EMIT connectionDeleted(gc);
        return true;
      });
      m_pwToQt.emplace(nodeInfo.id.value, ghostMatch);
      auto &data = m_nodes[ghostMatch];
      data.info = nodeInfo;

      auto portsResult = m_client->ListPorts(nodeInfo.id);
      if (portsResult.ok()) {
        data.inputPorts.clear();
        data.outputPorts.clear();
        for (const auto &port : portsResult.value) {
          if (port.is_input) {
            data.inputPorts.push_back(port);
          } else {
            data.outputPorts.push_back(port);
          }
        }
        std::sort(data.inputPorts.begin(), data.inputPorts.end(),
                  [](const auto &a, const auto &b) { return a.name < b.name; });
        std::sort(data.outputPorts.begin(), data.outputPorts.end(),
                  [](const auto &a, const auto &b) { return a.name < b.name; });
      }

      Q_EMIT nodeUpdated(ghostMatch);
      continue;
    }

    auto portsResult = m_client->ListPorts(nodeInfo.id);
    std::vector<warppipe::PortInfo> inputs;
    std::vector<warppipe::PortInfo> outputs;
    if (portsResult.ok()) {
      for (const auto &port : portsResult.value) {
        if (port.is_input) {
          inputs.push_back(port);
        } else {
          outputs.push_back(port);
        }
      }
      std::sort(inputs.begin(), inputs.end(),
                [](const auto &a, const auto &b) { return a.name < b.name; });
      std::sort(outputs.begin(), outputs.end(),
                [](const auto &a, const auto &b) { return a.name < b.name; });
    }

    QtNodes::NodeId qtId = newNodeId();
    WarpNodeData data;
    data.info = nodeInfo;
    data.inputPorts = std::move(inputs);
    data.outputPorts = std::move(outputs);

    auto [nodeIt, _] = m_nodes.emplace(qtId, std::move(data));
    m_pwToQt.emplace(nodeInfo.id.value, qtId);

    auto pendingIt = m_pendingPositions.find(nodeInfo.name);
    if (pendingIt != m_pendingPositions.end()) {
      m_positions.emplace(qtId, pendingIt->second);
      m_pendingPositions.erase(pendingIt);
    } else {
      auto savedIt = m_savedPositions.find(nodeInfo.name);
      if (savedIt != m_savedPositions.end()) {
        m_positions.emplace(qtId, savedIt->second);
      } else {
        m_positions.emplace(qtId, nextPosition(nodeIt->second));
      }
    }

    Q_EMIT nodeCreated(qtId);
  }

  std::vector<uint32_t> disappearedPwIds;
  for (const auto &[pwId, qtId] : m_pwToQt) {
    if (seenPwIds.find(pwId) == seenPwIds.end()) {
      disappearedPwIds.push_back(pwId);
    }
  }
  for (uint32_t pwId : disappearedPwIds) {
    auto it = m_pwToQt.find(pwId);
    if (it == m_pwToQt.end()) {
      continue;
    }
    QtNodes::NodeId qtId = it->second;
    auto nodeIt = m_nodes.find(qtId);
    if (nodeIt == m_nodes.end()) {
      continue;
    }
    WarpNodeType type = classifyNode(nodeIt->second.info);
    if (type == WarpNodeType::kApplication) {
      m_ghostNodes.insert(qtId);
      m_pwToQt.erase(it);
      Q_EMIT nodeUpdated(qtId);
    } else {
      m_pwToQt.erase(it);
      deleteNode(qtId);
    }
  }

  auto linksResult = m_client->ListLinks();
  if (linksResult.ok()) {
    std::unordered_set<uint32_t> seenLinkIds;
    for (const auto &link : linksResult.value) {
      seenLinkIds.insert(link.id.value);

      if (m_linkIdToConn.find(link.id.value) != m_linkIdToConn.end()) {
        continue;
      }

      auto outNodeIt = m_pwToQt.end();
      auto inNodeIt = m_pwToQt.end();
      QtNodes::PortIndex outPortIdx = 0;
      QtNodes::PortIndex inPortIdx = 0;
      bool found = false;

      for (const auto &[qtId, nodeData] : m_nodes) {
        for (size_t i = 0; i < nodeData.outputPorts.size(); ++i) {
          if (nodeData.outputPorts[i].id.value == link.output_port.value) {
            outNodeIt = m_pwToQt.find(nodeData.info.id.value);
            outPortIdx = static_cast<QtNodes::PortIndex>(i);
          }
        }
        for (size_t i = 0; i < nodeData.inputPorts.size(); ++i) {
          if (nodeData.inputPorts[i].id.value == link.input_port.value) {
            inNodeIt = m_pwToQt.find(nodeData.info.id.value);
            inPortIdx = static_cast<QtNodes::PortIndex>(i);
          }
        }
      }

      if (outNodeIt != m_pwToQt.end() && inNodeIt != m_pwToQt.end()) {
        found = true;
      }

      if (found) {
        QtNodes::ConnectionId connId{outNodeIt->second, outPortIdx,
                                     inNodeIt->second, inPortIdx};
        if (m_connections.find(connId) == m_connections.end()) {
          m_connections.insert(connId);
          m_linkIdToConn.emplace(link.id.value, connId);
          Q_EMIT connectionCreated(connId);
        }
      }
    }

    std::vector<uint32_t> staleLinkIds;
    for (const auto &[linkId, connId] : m_linkIdToConn) {
      if (seenLinkIds.find(linkId) == seenLinkIds.end()) {
        staleLinkIds.push_back(linkId);
      }
    }
    for (uint32_t linkId : staleLinkIds) {
      auto it = m_linkIdToConn.find(linkId);
      if (it != m_linkIdToConn.end()) {
        QtNodes::ConnectionId connId = it->second;
        bool outIsGhost =
            m_ghostNodes.find(connId.outNodeId) != m_ghostNodes.end();
        bool inIsGhost =
            m_ghostNodes.find(connId.inNodeId) != m_ghostNodes.end();

        if (outIsGhost || inIsGhost) {
          m_ghostConnections.insert(connId);
        } else {
          auto connIt = m_connections.find(connId);
          if (connIt != m_connections.end()) {
            m_connections.erase(connIt);
            Q_EMIT connectionDeleted(connId);
          }
        }
        m_linkIdToConn.erase(it);
      }
    }
  }

  if (!m_pendingGhostConnections.empty()) {
    auto it = m_pendingGhostConnections.begin();
    while (it != m_pendingGhostConnections.end()) {
      QtNodes::NodeId outQtId = 0;
      QtNodes::NodeId inQtId = 0;
      for (const auto &[qtId, data] : m_nodes) {
        if (data.info.name == it->outNodeName)
          outQtId = qtId;
        if (data.info.name == it->inNodeName)
          inQtId = qtId;
      }
      if (outQtId == 0 || inQtId == 0) {
        ++it;
        continue;
      }

      auto outNodeIt = m_nodes.find(outQtId);
      auto inNodeIt = m_nodes.find(inQtId);
      QtNodes::PortIndex outIdx = -1;
      QtNodes::PortIndex inIdx = -1;
      for (size_t i = 0; i < outNodeIt->second.outputPorts.size(); ++i) {
        if (outNodeIt->second.outputPorts[i].name == it->outPortName) {
          outIdx = static_cast<QtNodes::PortIndex>(i);
          break;
        }
      }
      for (size_t i = 0; i < inNodeIt->second.inputPorts.size(); ++i) {
        if (inNodeIt->second.inputPorts[i].name == it->inPortName) {
          inIdx = static_cast<QtNodes::PortIndex>(i);
          break;
        }
      }
      if (outIdx < 0 || inIdx < 0) {
        ++it;
        continue;
      }

      QtNodes::ConnectionId connId{outQtId, outIdx, inQtId, inIdx};
      if (m_connections.find(connId) == m_connections.end()) {
        m_connections.insert(connId);
        m_ghostConnections.insert(connId);
        Q_EMIT connectionCreated(connId);
      }
      it = m_pendingGhostConnections.erase(it);
    }
  }

  m_refreshing = false;
}

const WarpNodeData *
WarpGraphModel::warpNodeData(QtNodes::NodeId nodeId) const {
  auto it = m_nodes.find(nodeId);
  if (it != m_nodes.end()) {
    return &it->second;
  }
  return nullptr;
}

QtNodes::NodeId WarpGraphModel::qtNodeIdForPw(uint32_t pwNodeId) const {
  auto it = m_pwToQt.find(pwNodeId);
  if (it != m_pwToQt.end()) {
    return it->second;
  }
  return 0;
}

void WarpGraphModel::setPendingPosition(const std::string &nodeName,
                                        QPointF pos) {
  m_pendingPositions[nodeName] = pos;
}

QString WarpGraphModel::captionForNode(const warppipe::NodeInfo &info) {
  if (!info.description.empty()) {
    return QString::fromStdString(info.description);
  }
  if (!info.application_name.empty() && info.application_name != info.name) {
    return QString::fromStdString(info.application_name);
  }
  return QString::fromStdString(info.name);
}

QSize WarpGraphModel::estimateNodeSize(const WarpNodeData &data) {
  int maxPorts = static_cast<int>(
      std::max(data.inputPorts.size(), data.outputPorts.size()));
  int height = std::max(80, 50 + maxPorts * 28);

  QString caption = captionForNode(data.info);
  int captionWidth = caption.length() * 8 + 40;

  int maxInputLen = 0;
  int maxOutputLen = 0;
  for (const auto &p : data.inputPorts)
    maxInputLen = std::max(maxInputLen, static_cast<int>(p.name.length()));
  for (const auto &p : data.outputPorts)
    maxOutputLen = std::max(maxOutputLen, static_cast<int>(p.name.length()));
  int portWidth = (maxInputLen + maxOutputLen) * 7 + 60;

  int width = std::max(180, std::max(captionWidth, portWidth));
  return QSize(width, height);
}

QPointF WarpGraphModel::nextPosition(const WarpNodeData &data) {
  QSize size = estimateNodeSize(data);
  double nodeW = size.width();
  double nodeH = size.height();

  if (m_nextX + nodeW > kMaxRowWidth && m_nextX > 0) {
    m_nextX = 0.0;
    m_nextY += m_rowMaxHeight + kVerticalGap;
    m_rowMaxHeight = 0.0;
  }

  QPointF pos(m_nextX, m_nextY);
  m_nextX += nodeW + kHorizontalGap;
  m_rowMaxHeight = std::max(m_rowMaxHeight, nodeH);
  return pos;
}

bool WarpGraphModel::isGhost(QtNodes::NodeId nodeId) const {
  return m_ghostNodes.find(nodeId) != m_ghostNodes.end();
}

uint32_t WarpGraphModel::findPwNodeIdByName(const std::string &name) const {
  for (const auto &[qtId, data] : m_nodes) {
    if (data.info.name == name) {
      return data.info.id.value;
    }
  }
  return 0;
}

WarpNodeType
WarpGraphModel::classifyNode(const warppipe::NodeInfo &info) {
  const std::string &mc = info.media_class;

  if (mc == "Audio/Sink" || mc == "Audio/Duplex") {
    return info.is_virtual ? WarpNodeType::kVirtualSink
                           : WarpNodeType::kHardwareSink;
  }
  if (mc == "Audio/Source") {
    return info.is_virtual ? WarpNodeType::kVirtualSource
                           : WarpNodeType::kHardwareSource;
  }
  if (mc == "Stream/Output/Audio" || mc == "Stream/Input/Audio") {
    return WarpNodeType::kApplication;
  }

  return WarpNodeType::kUnknown;
}

void WarpGraphModel::saveLayout(const QString &path) const {
  ViewState vs{};
  saveLayout(path, vs);
}

void WarpGraphModel::saveLayout(const QString &path,
                                const ViewState &viewState) const {
  QJsonArray nodesArray;
  for (const auto &[qtId, data] : m_nodes) {
    auto posIt = m_positions.find(qtId);
    if (posIt == m_positions.end()) {
      continue;
    }
    QJsonObject nodeObj;
    nodeObj["name"] = QString::fromStdString(data.info.name);
    nodeObj["x"] = posIt->second.x();
    nodeObj["y"] = posIt->second.y();
    nodesArray.append(nodeObj);
  }

  QJsonArray ghostsArray;
  for (const auto &ghostId : m_ghostNodes) {
    auto nodeIt = m_nodes.find(ghostId);
    if (nodeIt == m_nodes.end()) {
      continue;
    }
    const auto &data = nodeIt->second;
    QJsonObject ghostObj;
    ghostObj["name"] = QString::fromStdString(data.info.name);
    ghostObj["description"] = QString::fromStdString(data.info.description);
    ghostObj["media_class"] = QString::fromStdString(data.info.media_class);
    ghostObj["application_name"] =
        QString::fromStdString(data.info.application_name);

    auto posIt = m_positions.find(ghostId);
    if (posIt != m_positions.end()) {
      ghostObj["x"] = posIt->second.x();
      ghostObj["y"] = posIt->second.y();
    }

    QJsonArray inPorts;
    for (const auto &port : data.inputPorts) {
      QJsonObject p;
      p["id"] = static_cast<int>(port.id.value);
      p["name"] = QString::fromStdString(port.name);
      inPorts.append(p);
    }
    ghostObj["input_ports"] = inPorts;

    QJsonArray outPorts;
    for (const auto &port : data.outputPorts) {
      QJsonObject p;
      p["id"] = static_cast<int>(port.id.value);
      p["name"] = QString::fromStdString(port.name);
      outPorts.append(p);
    }
    ghostObj["output_ports"] = outPorts;

    ghostsArray.append(ghostObj);
  }

  QJsonArray ghostConnsArray;
  for (const auto &conn : m_ghostConnections) {
    auto outIt = m_nodes.find(conn.outNodeId);
    auto inIt = m_nodes.find(conn.inNodeId);
    if (outIt == m_nodes.end() || inIt == m_nodes.end()) {
      continue;
    }
    auto outIdx = static_cast<size_t>(conn.outPortIndex);
    auto inIdx = static_cast<size_t>(conn.inPortIndex);
    if (outIdx >= outIt->second.outputPorts.size() ||
        inIdx >= inIt->second.inputPorts.size()) {
      continue;
    }
    QJsonObject connObj;
    connObj["out_node"] =
        QString::fromStdString(outIt->second.info.name);
    connObj["out_port"] =
        QString::fromStdString(outIt->second.outputPorts[outIdx].name);
    connObj["in_node"] =
        QString::fromStdString(inIt->second.info.name);
    connObj["in_port"] =
        QString::fromStdString(inIt->second.inputPorts[inIdx].name);
    ghostConnsArray.append(connObj);
  }

  QJsonObject root;
  root["version"] = 2;
  root["nodes"] = nodesArray;
  root["ghosts"] = ghostsArray;
  root["ghost_connections"] = ghostConnsArray;

  if (viewState.valid) {
    QJsonObject viewObj;
    viewObj["scale"] = viewState.scale;
    viewObj["center_x"] = viewState.centerX;
    viewObj["center_y"] = viewState.centerY;
    root["view"] = viewObj;
  }

  QFileInfo fi(path);
  QDir dir = fi.absoluteDir();
  if (!dir.exists()) {
    dir.mkpath(".");
  }

  QFile file(path);
  if (file.open(QIODevice::WriteOnly | QIODevice::Truncate)) {
    file.write(QJsonDocument(root).toJson(QJsonDocument::Compact));
  }
}

void WarpGraphModel::clearSavedPositions() {
  m_savedPositions.clear();
  m_positions.clear();
}

WarpGraphModel::ViewState WarpGraphModel::savedViewState() const {
  return m_savedViewState;
}

bool WarpGraphModel::loadLayout(const QString &path) {
  QFile file(path);
  if (!file.open(QIODevice::ReadOnly)) {
    return false;
  }

  QJsonDocument doc = QJsonDocument::fromJson(file.readAll());
  if (!doc.isObject()) {
    return false;
  }

  QJsonObject root = doc.object();
  int version = root["version"].toInt();
  if (version < 1 || version > 2) {
    return false;
  }

  m_savedPositions.clear();
  QJsonArray nodesArray = root["nodes"].toArray();
  for (const auto &val : nodesArray) {
    QJsonObject obj = val.toObject();
    std::string name = obj["name"].toString().toStdString();
    double x = obj["x"].toDouble();
    double y = obj["y"].toDouble();
    m_savedPositions[name] = QPointF(x, y);
  }

  m_savedViewState = {};
  if (root.contains("view")) {
    QJsonObject viewObj = root["view"].toObject();
    m_savedViewState.scale = viewObj["scale"].toDouble(1.0);
    m_savedViewState.centerX = viewObj["center_x"].toDouble();
    m_savedViewState.centerY = viewObj["center_y"].toDouble();
    m_savedViewState.valid = true;
  }

  if (root.contains("ghosts")) {
    QJsonArray ghostsArray = root["ghosts"].toArray();
    for (const auto &val : ghostsArray) {
      QJsonObject obj = val.toObject();
      std::string name = obj["name"].toString().toStdString();

      bool alreadyExists = false;
      for (const auto &[_, data] : m_nodes) {
        if (data.info.name == name) {
          alreadyExists = true;
          break;
        }
      }
      if (alreadyExists) {
        continue;
      }

      warppipe::NodeInfo info;
      info.id = warppipe::NodeId{0};
      info.name = name;
      info.description = obj["description"].toString().toStdString();
      info.media_class = obj["media_class"].toString().toStdString();
      info.application_name =
          obj["application_name"].toString().toStdString();

      WarpNodeData data;
      data.info = info;

      for (const auto &pval : obj["input_ports"].toArray()) {
        QJsonObject p = pval.toObject();
        warppipe::PortInfo port;
        port.id = warppipe::PortId{
            static_cast<uint32_t>(p["id"].toInt())};
        port.node = info.id;
        port.name = p["name"].toString().toStdString();
        port.is_input = true;
        data.inputPorts.push_back(port);
      }
      for (const auto &pval : obj["output_ports"].toArray()) {
        QJsonObject p = pval.toObject();
        warppipe::PortInfo port;
        port.id = warppipe::PortId{
            static_cast<uint32_t>(p["id"].toInt())};
        port.node = info.id;
        port.name = p["name"].toString().toStdString();
        port.is_input = false;
        data.outputPorts.push_back(port);
      }

      QtNodes::NodeId qtId = newNodeId();
      m_nodes.emplace(qtId, std::move(data));
      m_ghostNodes.insert(qtId);

      if (obj.contains("x") && obj.contains("y")) {
        m_positions.emplace(qtId, QPointF(obj["x"].toDouble(),
                                          obj["y"].toDouble()));
      }
      m_savedPositions[name] =
          m_positions.count(qtId)
              ? m_positions.at(qtId)
              : QPointF(0, 0);

      Q_EMIT nodeCreated(qtId);
    }
  }

  if (root.contains("ghost_connections")) {
    m_pendingGhostConnections.clear();
    QJsonArray gcArray = root["ghost_connections"].toArray();
    for (const auto &val : gcArray) {
      QJsonObject obj = val.toObject();
      PendingGhostConnection pgc;
      pgc.outNodeName = obj["out_node"].toString().toStdString();
      pgc.outPortName = obj["out_port"].toString().toStdString();
      pgc.inNodeName = obj["in_node"].toString().toStdString();
      pgc.inPortName = obj["in_port"].toString().toStdString();
      m_pendingGhostConnections.push_back(std::move(pgc));
    }
  }

  return !m_savedPositions.empty() || !m_ghostNodes.empty();
}

void WarpGraphModel::autoArrange() {
  struct Column {
    std::vector<QtNodes::NodeId> ids;
    double maxWidth = 0.0;
  };

  Column sources;
  Column apps;
  Column sinks;

  for (const auto &[qtId, data] : m_nodes) {
    WarpNodeType type = classifyNode(data.info);
    QSize sz = estimateNodeSize(data);
    double w = sz.width();

    switch (type) {
    case WarpNodeType::kHardwareSource:
    case WarpNodeType::kVirtualSource:
      sources.ids.push_back(qtId);
      sources.maxWidth = std::max(sources.maxWidth, w);
      break;
    case WarpNodeType::kApplication:
      apps.ids.push_back(qtId);
      apps.maxWidth = std::max(apps.maxWidth, w);
      break;
    default:
      sinks.ids.push_back(qtId);
      sinks.maxWidth = std::max(sinks.maxWidth, w);
      break;
    }
  }

  auto layoutColumn = [&](Column &col, double xOffset) {
    double y = 0.0;
    for (QtNodes::NodeId id : col.ids) {
      auto it = m_nodes.find(id);
      if (it == m_nodes.end()) continue;
      QSize sz = estimateNodeSize(it->second);
      m_positions[id] = QPointF(xOffset, y);
      Q_EMIT nodePositionUpdated(id);
      y += sz.height() + kVerticalGap;
    }
  };

  double x = 0.0;
  layoutColumn(sources, x);
  x += sources.maxWidth + kHorizontalGap * 3;
  layoutColumn(apps, x);
  x += apps.maxWidth + kHorizontalGap * 3;
  layoutColumn(sinks, x);
}

QVariant WarpGraphModel::styleForNode(WarpNodeType type, bool ghost) {
  QtNodes::NodeStyle style = QtNodes::StyleCollection::nodeStyle();

  QColor base;
  switch (type) {
  case WarpNodeType::kHardwareSink:
    base = QColor(72, 94, 118);
    break;
  case WarpNodeType::kHardwareSource:
    base = QColor(94, 72, 118);
    break;
  case WarpNodeType::kVirtualSink:
    base = QColor(62, 122, 104);
    break;
  case WarpNodeType::kVirtualSource:
    base = QColor(62, 104, 122);
    break;
  case WarpNodeType::kApplication:
    base = QColor(138, 104, 72);
    break;
  default:
    base = QColor(86, 94, 108);
    break;
  }

  if (ghost) {
    style.GradientColor0 = base.darker(150);
    style.GradientColor1 = base.darker(160);
    style.GradientColor2 = base.darker(170);
    style.GradientColor3 = base.darker(180);
    style.NormalBoundaryColor = base.darker(130);
    style.FontColor = QColor(160, 168, 182);
    style.FontColorFaded = QColor(120, 128, 142);
    style.ConnectionPointColor = QColor(140, 148, 160);
    style.FilledConnectionPointColor = QColor(180, 140, 80);
  } else {
    style.GradientColor0 = base.lighter(120);
    style.GradientColor1 = base.lighter(108);
    style.GradientColor2 = base.darker(105);
    style.GradientColor3 = base.darker(120);
    style.NormalBoundaryColor = base.lighter(135);
    style.FontColor = QColor(236, 240, 246);
    style.FontColorFaded = QColor(160, 168, 182);
    style.ConnectionPointColor = QColor(200, 208, 220);
    style.FilledConnectionPointColor = QColor(255, 165, 0);
  }

  style.Opacity = 1.0f;
  style.SelectedBoundaryColor = QColor(255, 165, 0);
  style.PenWidth = 1.3f;
  style.HoveredPenWidth = 2.4f;
  style.ConnectionPointDiameter = 10.0f;

  return style.toJson().toVariantMap();
}