// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "ZMQMetadataSocket.h"
#include "../frame_serialize/CBORStream2Serializer.h"
#include "../common/ZMQWrappers.h"
#include "../common/Definitions.h"

#define MESSAGE_SIZE_FOR_IMAGE (20*1024) // 20 kB


ZMQMetadataSocket::ZMQMetadataSocket(const std::string &addr)
    : metadata_message({}), socket(ZMQSocketType::Pub), counter(std::chrono::seconds(1)) {
    socket.SendWaterMark(DefaultSendWatermark);
    socket.Bind(addr);

    first_image = true;
}

void ZMQMetadataSocket::StartDataCollection(const StartMessage &message) {
    std::unique_lock ul(m);

    metadata_message = MetadataMessage{};
    if (buffer.size() < MESSAGE_SIZE_FOR_START_END)
        buffer.resize(MESSAGE_SIZE_FOR_START_END);
    CBORStream2Serializer serializer(buffer.data(), buffer.size());
    serializer.SerializeSequenceStart(message);
    socket.Send(buffer.data(), serializer.GetBufferSize(), true);
    first_image = true;
}

// Called without m: serialising a batch must not hold up the image threads
void ZMQMetadataSocket::Send(const MetadataMessage &batch) {
    size_t size = 0;
    for (const auto &i: batch.images)
        size += MESSAGE_SIZE_FOR_IMAGE + 16 * i.az_int_profile.size() + 9 * i.adu_histogram.size();

    std::vector<uint8_t> batch_buffer(size);
    CBORStream2Serializer serializer(batch_buffer.data(), batch_buffer.size());
    serializer.SerializeMetadata(batch);
    socket.Send(batch_buffer.data(), serializer.GetBufferSize(), false);
}

void ZMQMetadataSocket::AddDataMessage(const DataMessage &msg) {
    // Disabled: nothing is kept
    if (!counter.GetPeriod())
        return;

    // Spots, reflections and the image are not part of the metadata message
    DataMessage copy = msg;
    copy.spots = {};
    copy.reflections = {};
    copy.image = CompressedImage();

    MetadataMessage batch;
    {
        std::unique_lock ul(m);

        if (first_image) {
            // GeneratePreview() is designed to fire on the first try. This would result however in situation where
            // first metadata packet always contains only 1 image, and this is not desired behavior.
            // Therefore, on first use there is one extra "artificial" check of generate preview, that starts the clock
            counter.GeneratePreview();

            metadata_message.run_name = msg.run_name;
            metadata_message.run_number = msg.run_number;
            first_image = false;
        }

        metadata_message.images.emplace_back(std::move(copy));
        if (!counter.GeneratePreview())
            return;

        batch.run_name = metadata_message.run_name;
        batch.run_number = metadata_message.run_number;
        std::swap(batch.images, metadata_message.images);
    }
    Send(batch);
}

void ZMQMetadataSocket::EndDataCollection(const EndMessage &message) {
    std::unique_lock ul(m);
    if (!metadata_message.images.empty()) {
        Send(metadata_message);
        metadata_message.images.clear();
    }

    if (buffer.size() < MESSAGE_SIZE_FOR_START_END)
        buffer.resize(MESSAGE_SIZE_FOR_START_END);
    CBORStream2Serializer serializer(buffer.data(), buffer.size());
    serializer.SerializeSequenceEnd(message);
    socket.Send(buffer.data(), serializer.GetBufferSize(), true);
}

ZMQMetadataSocket &ZMQMetadataSocket::Period(const std::optional<std::chrono::microseconds> &period) {
    counter.Period(period);
    return *this;
}

ZMQMetadataSocket& ZMQMetadataSocket::ImportSettings(const ZMQMetadataSettings &settings) {
    std::unique_lock ul(m);
    Period(settings.period);
    return *this;
}

ZMQMetadataSettings ZMQMetadataSocket::GetSettings() {
    std::unique_lock ul(m);
    return ZMQMetadataSettings{
        .period = counter.GetPeriod(),
        .address = socket.GetEndpointName()
    };
}

std::string ZMQMetadataSocket::GetAddress() {
    return socket.GetEndpointName();
}