/*
 * SPDX-License-Identifier: GPL-3.0-or-later
 */
#include "ninep.h"

#include <QLocalSocket>

// 9P2000 message type codes (plan9port fcall.h)
enum {
    Tversion = 100,
    Rversion = 101,
    Tattach = 104,
    Rattach = 105,
    Rerror = 107,
    Twalk = 110,
    Rwalk = 111,
    Topen = 112,
    Ropen = 113,
    Tread = 116,
    Rread = 117,
    Twrite = 118,
    Rwrite = 119,
    Tclunk = 120,
    Rclunk = 121,
};

static constexpr uint16_t NOTAG = 0xFFFF;
static constexpr uint32_t NOFID = 0xFFFFFFFFu;
static constexpr uint16_t Tag = 1;
static constexpr int ConnectTimeoutMs = 5000;
static constexpr uint32_t IOHDRSZ_LOCAL = 24;

// Little-endian field appenders
static void put1(QByteArray &b, uint8_t v) { b.append(char(v)); }
static void put2(QByteArray &b, uint16_t v)
{
    b.append(char(v & 0xFF));
    b.append(char((v >> 8) & 0xFF));
}
static void put4(QByteArray &b, uint32_t v)
{
    for (int i = 0; i < 4; i++)
        b.append(char((v >> (8 * i)) & 0xFF));
}
static void put8(QByteArray &b, uint64_t v)
{
    for (int i = 0; i < 8; i++)
        b.append(char((v >> (8 * i)) & 0xFF));
}
static void putStr(QByteArray &b, const QByteArray &s)
{
    put2(b, uint16_t(s.size()));
    b.append(s);
}

// Little-endian field reader
namespace {
struct Reader {
    const char *p;
    int remaining;
    bool ok = true;
    static Reader of(const QByteArray &b) { return Reader{b.constData(), int(b.size())}; }
    uint8_t u1()
    {
        if (remaining < 1) { ok = false; return 0; }
        uint8_t v = uint8_t(*p);
        p++; remaining--;
        return v;
    }
    uint16_t u2()
    {
        if (remaining < 2) { ok = false; return 0; }
        uint16_t v = uint8_t(p[0]) | (uint16_t(uint8_t(p[1])) << 8);
        p += 2; remaining -= 2;
        return v;
    }
    uint32_t u4()
    {
        if (remaining < 4) { ok = false; return 0; }
        uint32_t v = 0;
        for (int i = 0; i < 4; i++)
            v |= uint32_t(uint8_t(p[i])) << (8 * i);
        p += 4; remaining -= 4;
        return v;
    }
    QByteArray str()
    {
        uint16_t n = u2();
        if (!ok || remaining < int(n)) { ok = false; return QByteArray(); }
        QByteArray s(p, n);
        p += n; remaining -= n;
        return s;
    }
};
} // namespace

NineP::NineP() = default;

NineP::~NineP()
{
    close();
}

bool NineP::isConnected() const
{
    return m_sock && m_sock->state() == QLocalSocket::ConnectedState;
}

void NineP::close()
{
    if (m_sock) {
        if (m_sock->state() == QLocalSocket::ConnectedState)
            clunk(RootFid);
        m_sock->abort();
        delete m_sock;
        m_sock = nullptr;
    }
}

bool NineP::readExactly(char *buf, int n)
{
    static constexpr int ReadChunkTimeoutMs = 2000;
    int got = 0;
    while (got < n) {
        if (m_sock->bytesAvailable() == 0) {
            if (!m_sock->waitForReadyRead(ReadChunkTimeoutMs)) {
                m_error = QStringLiteral("socket read failed: %1").arg(m_sock->errorString());
                return false;
            }
        }
        const qint64 r = m_sock->read(buf + got, n - got);
        if (r < 0) {
            m_error = QStringLiteral("socket read error: %1").arg(m_sock->errorString());
            return false;
        }
        got += int(r);
    }
    return true;
}

bool NineP::sendMsg(uint8_t type, uint16_t tag, const QByteArray &body)
{
    QByteArray msg;
    const uint32_t size = 4 + 1 + 2 + uint32_t(body.size());
    put4(msg, size);
    put1(msg, type);
    put2(msg, tag);
    msg.append(body);
    if (m_sock->write(msg) != msg.size()) {
        m_error = QStringLiteral("socket write failed: %1").arg(m_sock->errorString());
        return false;
    }
    if (!m_sock->waitForBytesWritten(ConnectTimeoutMs)) {
        m_error = QStringLiteral("socket flush timed out");
        return false;
    }
    return true;
}

bool NineP::recvMsg(uint8_t *type, uint16_t *tag, QByteArray *body)
{
    char hdr[7];
    if (!readExactly(hdr, 7))
        return false;
    const uint32_t size = uint32_t(uint8_t(hdr[0])) | (uint32_t(uint8_t(hdr[1])) << 8)
        | (uint32_t(uint8_t(hdr[2])) << 16) | (uint32_t(uint8_t(hdr[3])) << 24);
    if (size < 7) {
        m_error = QStringLiteral("9P message too short (%1)").arg(size);
        return false;
    }
    *type = uint8_t(hdr[4]);
    *tag = uint16_t(uint8_t(hdr[5]) | (uint16_t(uint8_t(hdr[6])) << 8));
    const int bodyLen = int(size) - 7;
    QByteArray b;
    b.resize(bodyLen);
    if (bodyLen > 0 && !readExactly(b.data(), bodyLen))
        return false;
    if (*type == Rerror) {
        Reader r = Reader::of(b);
        const QByteArray ename = r.str();
        m_error = QStringLiteral("9P Rerror: %1").arg(QString::fromUtf8(ename));
        return false;
    }
    *body = b;
    return true;
}

bool NineP::connectAndAttach(const QString &socketPath, const QString &user,
                             const QString &aname)
{
    close();
    m_sock = new QLocalSocket();
    m_sock->connectToServer(socketPath, QLocalSocket::ReadWrite);
    if (!m_sock->waitForConnected(ConnectTimeoutMs)) {
        m_error = QStringLiteral("connect %1: %2").arg(socketPath, m_sock->errorString());
        delete m_sock;
        m_sock = nullptr;
        return false;
    }

    // Tversion: msize[4] version[s]
    {
        QByteArray body;
        put4(body, m_msize);
        putStr(body, QByteArrayLiteral("9P2000"));
        if (!sendMsg(Tversion, NOTAG, body)) {
            close();
            return false;
        }
        uint8_t type;
        uint16_t tag;
        QByteArray rb;
        if (!recvMsg(&type, &tag, &rb) || type != Rversion) {
            if (m_error.isEmpty())
                m_error = QStringLiteral("unexpected reply to Tversion");
            close();
            return false;
        }
        Reader r = Reader::of(rb);
        const uint32_t srvMsize = r.u4();
        const QByteArray srvVer = r.str();
        if (!r.ok || !srvVer.startsWith("9P2000")) {
            m_error = QStringLiteral("server declined 9P2000 (got '%1')")
                          .arg(QString::fromUtf8(srvVer));
            close();
            return false;
        }
        if (srvMsize > 0 && srvMsize < m_msize)
            m_msize = srvMsize;
    }

    // Tattach: fid[4] afid[4] uname[s] aname[s]
    {
        QByteArray body;
        put4(body, RootFid);
        put4(body, NOFID);
        putStr(body, user.toUtf8());
        putStr(body, aname.toUtf8());
        if (!sendMsg(Tattach, Tag, body)) {
            close();
            return false;
        }
        uint8_t type;
        uint16_t tag;
        QByteArray rb;
        if (!recvMsg(&type, &tag, &rb) || type != Rattach) {
            if (m_error.isEmpty())
                m_error = QStringLiteral("unexpected reply to Tattach");
            close();
            return false;
        }
    }
    return true;
}

bool NineP::walk(uint32_t newFid, const QString &name)
{
    // Twalk: fid[4] newfid[4] nwname[2] nwname*(wname[s])
    QByteArray body;
    put4(body, RootFid);
    put4(body, newFid);
    const QByteArray n = name.toUtf8();
    put2(body, 1);
    putStr(body, n);
    if (!sendMsg(Twalk, Tag, body))
        return false;
    uint8_t type;
    uint16_t tag;
    QByteArray rb;
    if (!recvMsg(&type, &tag, &rb) || type != Rwalk) {
        if (m_error.isEmpty())
            m_error = QStringLiteral("walk '%1' failed").arg(name);
        return false;
    }
    Reader r = Reader::of(rb);
    const uint16_t nwqid = r.u2();
    if (!r.ok || nwqid != 1) {
        m_error = QStringLiteral("walk '%1' did not resolve").arg(name);
        return false;
    }
    return true;
}

bool NineP::open(uint32_t fid, uint8_t mode)
{
    QByteArray body;
    put4(body, fid);
    put1(body, mode);
    if (!sendMsg(Topen, Tag, body))
        return false;
    uint8_t type;
    uint16_t tag;
    QByteArray rb;
    if (!recvMsg(&type, &tag, &rb) || type != Ropen) {
        if (m_error.isEmpty())
            m_error = QStringLiteral("open failed");
        return false;
    }
    return true;
}

int NineP::write(uint32_t fid, uint64_t offset, const QByteArray &data)
{
    if (uint32_t(data.size()) + 4 + 1 + 2 + 4 + 8 + 4 > m_msize) {
        m_error = QStringLiteral("write exceeds msize");
        return -1;
    }
    QByteArray body;
    put4(body, fid);
    put8(body, offset);
    put4(body, uint32_t(data.size()));
    body.append(data);
    if (!sendMsg(Twrite, Tag, body))
        return -1;
    uint8_t type;
    uint16_t tag;
    QByteArray rb;
    if (!recvMsg(&type, &tag, &rb) || type != Rwrite) {
        if (m_error.isEmpty())
            m_error = QStringLiteral("write failed");
        return -1;
    }
    Reader r = Reader::of(rb);
    const uint32_t count = r.u4();
    if (!r.ok) {
        m_error = QStringLiteral("malformed Rwrite");
        return -1;
    }
    return int(count);
}

int NineP::read(uint32_t fid, uint64_t offset, uint32_t count, QByteArray *out)
{
    if (count > m_msize - IOHDRSZ_LOCAL)
        count = m_msize - IOHDRSZ_LOCAL;
    QByteArray body;
    put4(body, fid);
    put8(body, offset);
    put4(body, count);
    if (!sendMsg(Tread, Tag, body))
        return -1;
    uint8_t type;
    uint16_t tag;
    QByteArray rb;
    if (!recvMsg(&type, &tag, &rb) || type != Rread) {
        if (m_error.isEmpty())
            m_error = QStringLiteral("read failed");
        return -1;
    }
    Reader r = Reader::of(rb);
    const uint32_t n = r.u4();
    if (!r.ok || int(n) > rb.size() - 4) {
        m_error = QStringLiteral("malformed Rread");
        return -1;
    }
    if (out)
        *out = rb.mid(4, int(n));
    return int(n);
}

bool NineP::beginRead(uint32_t fid, uint64_t offset, uint32_t count)
{
    if (count > m_msize - IOHDRSZ_LOCAL)
        count = m_msize - IOHDRSZ_LOCAL;
    QByteArray body;
    put4(body, fid);
    put8(body, offset);
    put4(body, count);
    return sendMsg(Tread, Tag, body);
}

int NineP::recvReadReply(QByteArray *out, int firstByteTimeoutMs, bool *timedOut)
{
    if (timedOut)
        *timedOut = false;
    if (m_sock->bytesAvailable() == 0) {
        if (!m_sock->waitForReadyRead(firstByteTimeoutMs)) {
            if (m_sock->state() == QLocalSocket::ConnectedState) {
                if (timedOut)
                    *timedOut = true;
                return 0;
            }
            m_error = QStringLiteral("socket read failed: %1").arg(m_sock->errorString());
            return -1;
        }
    }
    uint8_t type;
    uint16_t tag;
    QByteArray rb;
    if (!recvMsg(&type, &tag, &rb) || type != Rread) {
        if (m_error.isEmpty())
            m_error = QStringLiteral("read reply failed");
        return -1;
    }
    Reader r = Reader::of(rb);
    const uint32_t n = r.u4();
    if (!r.ok || int(n) > rb.size() - 4) {
        m_error = QStringLiteral("malformed Rread");
        return -1;
    }
    if (out)
        *out = rb.mid(4, int(n));
    return int(n);
}

bool NineP::clunk(uint32_t fid)
{
    QByteArray body;
    put4(body, fid);
    if (!sendMsg(Tclunk, Tag, body))
        return false;
    uint8_t type;
    uint16_t tag;
    QByteArray rb;
    if (!recvMsg(&type, &tag, &rb) || type != Rclunk)
        return false;
    return true;
}