#include "WarpGraphModel.h"

#include <QJsonObject>
#include <QVariant>

#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;
  }
  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;
    }
    int maxPorts = static_cast<int>(
        std::max(data.inputPorts.size(), data.outputPorts.size()));
    int height = std::max(80, 50 + maxPorts * 28);
    return QSize(200, height);
  }
  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");
  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;
  }
  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;
  }

  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);

    auto existing = m_pwToQt.find(nodeInfo.id.value);
    if (existing != m_pwToQt.end()) {
      auto &data = m_nodes[existing->second];
      data.info = nodeInfo;
      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);

    m_nodes.emplace(qtId, std::move(data));
    m_pwToQt.emplace(nodeInfo.id.value, qtId);
    m_positions.emplace(qtId, nextPosition());

    Q_EMIT nodeCreated(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()) {
        auto connIt = m_connections.find(it->second);
        if (connIt != m_connections.end()) {
          QtNodes::ConnectionId connId = it->second;
          m_connections.erase(connIt);
          Q_EMIT connectionDeleted(connId);
        }
        m_linkIdToConn.erase(it);
      }
    }
  }
}

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;
}

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

QPointF WarpGraphModel::nextPosition() const {
  int count = static_cast<int>(m_nodes.size());
  double x = (count % 4) * 280.0;
  double y = (count / 4) * 200.0;
  return QPointF(x, y);
}