#ifndef CF5E5AFF_F218_4A25_95DF_8097D7C5685B
#define CF5E5AFF_F218_4A25_95DF_8097D7C5685B

#include "mana_api.hpp"
#include "mnode_spdlog.hpp"
#include <mesytec-mnode/mnode_cpp_types.h>
#include <mesytec-mnode/mnode_nng.h>

namespace mesytec::mnode::mana
{

struct MessageHeader
{
    enum MessageType
    {
        Invalid = 0,
        BeginRun = 42,
        EndRun = 43,
        EventData = 44,
    };

    u32 messageType;
};

struct EventDataHeader
{
    u16 eventIndex;
    u16 arrayCount;
    u32 sizeBytes;
};

// clang-format: off
// EventData message layout: MessageHeader, repeated(EventDataHeader +
// mana_offset_array_t[arrayCount])
// clang-format: on

class NngServerSink: public IManaSink
{
  public:
    explicit NngServerSink(nng_socket socket)
        : socket_(socket)
        , msg_(nng::make_unique_msg())
        , logger_(get_spdlog_logger("mana::NngServerSink"))
    {
    }

    void begin_run(const char *descriptor_json) override
    {
        auto len = std::strlen(descriptor_json);
        auto required = sizeof(MessageHeader) + len + 1;
        auto msg = nng::allocate_reserve_message(required);
        MessageHeader header{MessageHeader::BeginRun};
        nng_msg_append(msg.get(), &header, sizeof(header));
        nng_msg_append(msg.get(), descriptor_json, len + 1);

        logger_->info("Sending BeginRun message (size={})", nng_msg_len(msg.get()));

        int res = 0;

        do
        {
            if ((res = nng::send_message(socket_, msg)))
            {
                if (res != NNG_ETIMEDOUT)
                {
                    logger_->error("Error sending BeginRun message: {}", nng_strerror(res));
                    return;
                }
            }
        }
        while (res == NNG_ETIMEDOUT);

        logger_->debug("Waiting for response from client");

        while (true)
        {
            auto [msg, res] = nng::receive_message(socket_);

            if (res && res != NNG_ETIMEDOUT)
            {
                logger_->error("Error receiving response from client: {}", nng_strerror(res));
                return;
            }
            else if (!res)
            {
                logger_->info("Received response to BeginRun from client (size={})",
                              nng_msg_len(msg.get()));
                break;
            }
        }
    }

    void end_run(const char *descriptor_json) override
    {
        if (msg_)
        {
            if (!flush())
                return;
        }

        auto len = std::strlen(descriptor_json);
        auto required = sizeof(MessageHeader) + len + 1;
        auto msg = nng::allocate_reserve_message(required);
        MessageHeader header{MessageHeader::EndRun};
        nng_msg_append(msg.get(), &header, sizeof(header));
        nng_msg_append(msg.get(), descriptor_json, len + 1);
        logger_->info("Sending EndRun message (size={})", nng_msg_len(msg.get()));
        int res = 0;

        do
        {
            if ((res = nng::send_message(socket_, msg)))
            {
                if (res != NNG_ETIMEDOUT)
                {
                    logger_->error("Error sending EndRun message: {}", nng_strerror(res));
                    return;
                }
            }
        }
        while (res == NNG_ETIMEDOUT);

        logger_->info("Sent EndRun message");
    }

    void process_event(uint16_t eventIndex, mana_offset_array_t *arrays, size_t arrayCount,
                       size_t totalBytes) override
    {
        if (!msg_)
        {
            msg_ = nng::allocate_reserve_message(InitialMessageReserve);
            MessageHeader header{MessageHeader::EventData};
            nng_msg_append(msg_.get(), &header, sizeof(header));
        }

        const size_t required = sizeof(EventDataHeader) + totalBytes;

        if (nng::allocated_free_space(msg_.get()) < required &&
            nng_msg_len(msg_.get()) > sizeof(MessageHeader))
        {
            if (!flush())
                return;

            assert(!msg_);
            msg_ = nng::allocate_reserve_message(InitialMessageReserve);
            MessageHeader header{MessageHeader::EventData};
            nng_msg_append(msg_.get(), &header, sizeof(header));
        }

        // second check in case the event is larger than the initial reserve
        if (nng::allocated_free_space(msg_.get()) < required)
            nng_msg_reserve(msg_.get(), required);

        assert(nng::allocated_free_space(msg_.get()) >= required);
        assert(nng_msg_len(msg_.get()) >= sizeof(MessageHeader));

        EventDataHeader header{};
        header.eventIndex = eventIndex;
        header.arrayCount = arrayCount;
        header.sizeBytes = totalBytes;
        nng_msg_append(msg_.get(), &header, sizeof(header));
        nng_msg_append(msg_.get(), arrays, totalBytes);
    }

    void process_system_event(const uint32_t *data, size_t size) override
    {
        (void)data;
        (void)size;
    }

  private:
    bool flush()
    {
        if (!msg_)
            return false;

        logger_->trace("Sending EventData message of size {}", nng_msg_len(msg_.get()));

        while (true)
        {
            if (auto res = nng::send_message(socket_, msg_))
            {
                if (res != NNG_ETIMEDOUT)
                {
                    logger_->error("Error sending EventData message: {}", nng_strerror(res));
                    return false;
                }
            }
            else
                break;
        }

        return true;
    }

    static constexpr size_t InitialMessageReserve = 2 * 1024 * 1024;
    nng_socket socket_;
    nng::unique_msg msg_;
    std::shared_ptr<spdlog::logger> logger_;
};

inline void nng_client_run(nng_socket socket_, mana::IManaSink *sink, std::atomic<bool> &quit)
{
    auto receive_message = [socket_, &quit]()
    {
        auto rp = [&quit](int res) { return !quit && res == NNG_ETIMEDOUT; };
        return nng::receive_message_retry(socket_, 0, rp);
    };

    enum class State
    {
        BeginRun,
        Run,
    };

    auto logger = get_spdlog_logger("mana::nng_client");

    auto process_event_data = [&logger](nng_msg *msg, mana::IManaSink *sink)
    {
        while (nng_msg_len(msg) >= sizeof(mana::EventDataHeader))
        {
            auto header = nng::msg_trim_read<mana::EventDataHeader>(msg);
            assert(header);

            if (nng_msg_len(msg) < header->sizeBytes)
            {
                logger->error("Error reading event data: expected {} bytes, got {}",
                              header->sizeBytes, nng_msg_len(msg));
                return false;
            }

            auto arrays = reinterpret_cast<mana_offset_array_t *>(nng_msg_body(msg));
            sink->process_event(header->eventIndex, arrays, header->arrayCount, header->sizeBytes);
            nng_msg_trim(msg, header->sizeBytes);
        }

        return nng_msg_len(msg) == 0;
    };

    auto state = State::BeginRun;

    while (!quit)
    {
        switch (state)
        {
        case State::BeginRun:
        {
            logger->info("Waiting for BeginRun from server");

            auto [msg, res] = receive_message();

            if (res)
            {
                logger->error("Error receiving BeginRun message: {}", nng_strerror(res));
                return;
            }

            const size_t msgSize = nng_msg_len(msg.get());
            logger->info("State::BeginRun: Received message of size {}", msgSize);

            auto header = nng::msg_trim_read<mana::MessageHeader>(msg.get());

            if (!header)
            {
                logger->error("Error reading message header");
                return;
            }

            logger->info("State::BeginRun: messageType={}", static_cast<int>(header->messageType));

            if (header->messageType == mana::MessageHeader::BeginRun)
            {
                logger->info("Received BeginRun message (size={})", msgSize);
                std::string descriptor_json;
                std::copy(reinterpret_cast<const char *>(nng_msg_body(msg.get())),
                          reinterpret_cast<const char *>(nng_msg_body(msg.get())) + msgSize,
                          std::back_inserter(descriptor_json));

                sink->begin_run(descriptor_json.c_str());

                if (auto res = nng::send_empty_message(socket_))
                {
                    logger->error("Error sending response to BeginRun message: {}",
                                  nng_strerror(res));
                    return;
                }
                logger->info("Receiving Data...");
                state = State::Run;
            }
            else
            {
                logger->error("Error reading BeginRun header: messageType={}",
                              static_cast<int>(header->messageType));
                return;
            }
        }
        break;
        case State::Run:
        {
            auto [msg, res] = receive_message();

            if (res)
            {
                logger->error("State::Run: Error receiving message: {}", nng_strerror(res));
                return;
            }

            const size_t msgSize = nng_msg_len(msg.get());
            logger->trace("State::Run: Received message of size {}", msgSize);

            auto header = nng::msg_trim_read<mana::MessageHeader>(msg.get());

            if (!header)
            {
                logger->error("State::Run: Error reading message header, message too short");
                return;
            }

            logger->trace("State::Run: received message: messageType={}",
                          static_cast<int>(header->messageType));

            if (header->messageType == mana::MessageHeader::EndRun)
            {
                logger->info("State::Run: Received EndRun message");
                std::string descriptor_json;
                std::copy(reinterpret_cast<const char *>(nng_msg_body(msg.get())),
                          reinterpret_cast<const char *>(nng_msg_body(msg.get())) + msgSize,
                          std::back_inserter(descriptor_json));
                sink->end_run(descriptor_json.c_str());
                return;
            }
            else if (header->messageType == mana::MessageHeader::EventData)
            {
                if (!process_event_data(msg.get(), sink))
                    return;
            }
            else
            {
                logger->error("State::Run: Received unknown message type: {}",
                              static_cast<int>(header->messageType));
                return;
            }
        }
        break;
        }
    }
}

} // namespace mesytec::mnode::mana

#endif /* CF5E5AFF_F218_4A25_95DF_8097D7C5685B */
