#include "SwayIpcClient.hpp"

#include <QJsonArray>
#include <QJsonDocument>
#include <QJsonObject>
#include <QtGlobal>

#include <cstring>
#include <utility>

namespace {

constexpr char kMagic[] = "i3-ipc";
constexpr int kMagicSize = 6;
constexpr int kHeaderSize = 14;
constexpr int kDefaultTimeoutMs = 2000;

} // namespace

SwayIpcClient::SwayIpcClient(QObject *parent)
    : QObject(parent)
    , m_command(new QLocalSocket(this))
    , m_event(new QLocalSocket(this))
{
    connect(m_command, &QLocalSocket::disconnected, this, &SwayIpcClient::handleCommandDisconnected);
    connect(m_event, &QLocalSocket::disconnected, this, &SwayIpcClient::handleEventDisconnected);
    connect(m_event, &QLocalSocket::readyRead, this, &SwayIpcClient::processEventSocket);
}

void SwayIpcClient::processEventSocket()
{
    appendIncoming(m_event, &m_eventBuffer);
    Message message;
    while (takeMessage(&m_eventBuffer, &message)) {
        if (m_inRequest > 0)
            m_pendingEvents.append(std::move(message));
        else
            emit eventReceived(message.type, QJsonDocument::fromJson(message.payload));
    }
}

SwayIpcClient::~SwayIpcClient()
{
    disconnectFromSway();
}

QString SwayIpcClient::socketPath()
{
    return QString::fromLocal8Bit(qgetenv("SWAYSOCK"));
}

bool SwayIpcClient::isConnected() const
{
    return m_connected
        && m_command->state() == QLocalSocket::ConnectedState
        && m_event->state() == QLocalSocket::ConnectedState;
}

void SwayIpcClient::disconnectFromSway()
{
    m_event->blockSignals(true);
    m_command->blockSignals(true);
    if (m_event->state() != QLocalSocket::UnconnectedState)
        m_event->disconnectFromServer();
    if (m_command->state() != QLocalSocket::UnconnectedState)
        m_command->disconnectFromServer();
    m_event->blockSignals(false);
    m_command->blockSignals(false);
    m_eventBuffer.clear();
    m_pendingEvents.clear();

    if (m_connected) {
        m_connected = false;
        emit connectedChanged();
        emit disconnected();
    }
}

bool SwayIpcClient::connectToSway()
{
    const QString path = socketPath();
    if (path.isEmpty())
        return false;

    disconnectFromSway();

    m_command->connectToServer(path);
    if (!m_command->waitForConnected(kDefaultTimeoutMs)) {
        qWarning("nebula: Sway IPC command socket failed: %s", qPrintable(m_command->errorString()));
        return false;
    }

    m_event->connectToServer(path);
    if (!m_event->waitForConnected(kDefaultTimeoutMs)) {
        qWarning("nebula: Sway IPC event socket failed: %s", qPrintable(m_event->errorString()));
        m_command->disconnectFromServer();
        return false;
    }

    m_connected = true;
    emit connectedChanged();
    return true;
}

QJsonDocument SwayIpcClient::request(quint32 type, const QByteArray &payload)
{
    if (m_command->state() != QLocalSocket::ConnectedState)
        return {};

    ++m_inRequest;
    const bool wrote = writeMessage(m_command, type, payload);
    Message reply;
    const bool read = wrote && readMessage(m_command, &reply, kDefaultTimeoutMs);
    --m_inRequest;

    flushPendingEvents();

    if (!read) {
        qWarning("nebula: Sway IPC request failed (type %u)", type);
        return {};
    }

    return QJsonDocument::fromJson(reply.payload);
}

bool SwayIpcClient::subscribe(const QStringList &events)
{
    if (m_event->state() != QLocalSocket::ConnectedState)
        return false;

    QJsonArray array;
    for (const QString &event : events)
        array.append(event);

    // Subscribe on the event connection so unsolicited events never interleave
    // with command replies on the request socket.
    m_event->blockSignals(true);
    const bool wrote = writeMessage(m_event, Subscribe, QJsonDocument(array).toJson(QJsonDocument::Compact));
    Message reply;
    const bool read = wrote && readMessage(m_event, &reply, kDefaultTimeoutMs);
    m_event->blockSignals(false);

    if (m_event->bytesAvailable() > 0)
        processEventSocket();

    if (!read) {
        qWarning("nebula: Sway IPC subscribe failed");
        return false;
    }

    const QJsonDocument doc = QJsonDocument::fromJson(reply.payload);
    if (!doc.isObject() || !doc.object().value(QStringLiteral("success")).toBool()) {
        qWarning("nebula: Sway IPC subscribe was rejected");
        return false;
    }
    return true;
}

bool SwayIpcClient::writeMessage(QLocalSocket *socket, quint32 type, const QByteArray &payload)
{
    QByteArray header;
    header.resize(kHeaderSize);
    memcpy(header.data(), kMagic, kMagicSize);
    const quint32 size = static_cast<quint32>(payload.size());
    memcpy(header.data() + kMagicSize, &size, sizeof(size));
    memcpy(header.data() + kMagicSize + sizeof(size), &type, sizeof(type));

    if (socket->write(header) != header.size())
        return false;
    if (!payload.isEmpty() && socket->write(payload) != payload.size())
        return false;
    return socket->waitForBytesWritten(kDefaultTimeoutMs);
}

bool SwayIpcClient::readMessage(QLocalSocket *socket, Message *message, int timeoutMs)
{
    QByteArray buffer;
    while (!takeMessage(&buffer, message)) {
        if (socket->bytesAvailable() <= 0 && !socket->waitForReadyRead(timeoutMs))
            return false;
        appendIncoming(socket, &buffer);
    }
    return true;
}

void SwayIpcClient::appendIncoming(QLocalSocket *socket, QByteArray *buffer)
{
    buffer->append(socket->readAll());
}

bool SwayIpcClient::takeMessage(QByteArray *buffer, Message *message)
{
    if (buffer->size() < kHeaderSize)
        return false;
    if (!buffer->startsWith(kMagic)) {
        qWarning("nebula: Sway IPC header magic mismatch");
        buffer->clear();
        return false;
    }

    quint32 size = 0;
    quint32 type = 0;
    memcpy(&size, buffer->constData() + kMagicSize, sizeof(size));
    memcpy(&type, buffer->constData() + kMagicSize + sizeof(size), sizeof(type));

    const int total = kHeaderSize + static_cast<int>(size);
    if (buffer->size() < total)
        return false;

    message->type = type;
    message->payload = buffer->mid(kHeaderSize, static_cast<int>(size));
    buffer->remove(0, total);
    return true;
}

void SwayIpcClient::flushPendingEvents()
{
    if (m_inRequest > 0 || m_pendingEvents.isEmpty())
        return;

    const QVector<Message> pending = std::move(m_pendingEvents);
    m_pendingEvents.clear();
    for (const Message &message : pending)
        emit eventReceived(message.type, QJsonDocument::fromJson(message.payload));
}

void SwayIpcClient::handleCommandDisconnected()
{
    if (!m_connected)
        return;
    qWarning("nebula: Sway IPC command socket disconnected");
    m_connected = false;
    emit connectedChanged();
    emit disconnected();
}

void SwayIpcClient::handleEventDisconnected()
{
    if (!m_connected)
        return;
    qWarning("nebula: Sway IPC event socket disconnected");
    m_connected = false;
    emit connectedChanged();
    emit disconnected();
}