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

#include "HDF5FilePusher.h"
#include "../frame_serialize/CBORStream2Deserializer.h"
#include "../frame_serialize/CBORStream2Serializer.h"

HDF5FilePusher::HDF5FilePusher(const std::string &repub_address,
                               const std::optional<int32_t> &repub_watermark) {
    if (!repub_address.empty()) {
        repub_socket = std::make_unique<ZMQSocket>(ZMQSocketType::Push);
        repub_socket->SendWaterMark(repub_watermark.value_or(default_repub_watermark));
        repub_socket->SendTimeout(RepubTimeout);
        repub_socket->Bind(repub_address);
    }
}

void HDF5FilePusher::StartDataCollection(StartMessage &message) {
    if (writer)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher is already writing images");
    writer = std::make_unique<FileWriter>(message);
    writer_future = std::async(std::launch::async, &HDF5FilePusher::WriterThread, this);
    images_written = 0;

    if (repub_socket) {
        try {
            StartMessage repub_message = message;
            repub_message.writer_notification_zmq_addr = "";

            size_t approx_size = 1024 * 1024;
            for (const auto &[key, value] : repub_message.pixel_mask)
                approx_size += value.size() * sizeof(uint32_t);

            std::vector<uint8_t> serialization_buffer(approx_size);
            CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size());
            serializer.SerializeSequenceStart(repub_message);

            repub_active = repub_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
            if (repub_active)
                logger.Info("Republish active");
        } catch (const JFJochException &e) {
            logger.ErrorException(e);
        }
    }
}

bool HDF5FilePusher::EndDataCollection(const EndMessage &message) {
    if (writer_future.valid()) {
        writer_queue.PutBlocking({.end = true});
        writer_future.get();
    } else
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher wasn't writing files");

    if (!writer)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher not ready for writing images");
    writer->WriteHDF5(message);
    writer->Finalize();
    writer.reset();

    if (repub_socket) {
        try {
            size_t approx_size = 1024 * 1024;
            std::vector<uint8_t> serialization_buffer(approx_size);
            CBORStream2Serializer serializer(serialization_buffer.data(), serialization_buffer.size());
            serializer.SerializeSequenceEnd(message);

            if (repub_active)
                repub_socket->Send(serialization_buffer.data(), serializer.GetBufferSize(), true);
        } catch (const JFJochException &e) {
            logger.ErrorException(e);
        }

        if (repub_active)
            logger.Info("Republish finished");
        repub_active = false;
    }

    return true;
}

bool HDF5FilePusher::SendImage(const uint8_t *image_data, size_t image_size, int64_t image_number) {
    if (!writer)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher not ready for sending");

    auto deserialized = CBORStream2Deserialize(image_data, image_size);
    if (deserialized->data_message) {
        writer->Write(*deserialized->data_message);
        images_written++;

        if (repub_socket && repub_active) {
            try {
                repub_socket->Send(image_data, image_size, false);
            } catch (const JFJochException &e) {
                logger.ErrorException(e);
            }
        }
    } else
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "HDF5FilePusher::SendImage accepts only data image");
    return true;
}

void HDF5FilePusher::SendImage(ZeroCopyReturnValue &z) {
    writer_queue.PutBlocking(ImagePusherQueueElement{
        .image_data = (uint8_t *) z.GetImage(),
        .z = &z,
        .end = false
    });
}

bool HDF5FilePusher::SendCalibration(const CompressedImage &message) {
    if (!writer)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Image pusher not ready for sending");

    writer->WriteHDF5(message);
    return true;
}

void HDF5FilePusher::WriterThread() {
    ImagePusherQueueElement e = writer_queue.GetBlocking();
    while (!e.end) {
        SendImage(e.image_data, e.z->GetImageSize(), e.z->GetImageNumber());
        e.z->release();
        e = writer_queue.GetBlocking();
    }
}

std::string HDF5FilePusher::PrintSetup() const {
    std::filesystem::path currentPath = std::filesystem::current_path();
    return "HDF5FilePusher: Images are written directly to file in base directory " + currentPath.string();
}

std::optional<uint64_t> HDF5FilePusher::GetImagesWritten() const {
    return images_written;
}

size_t HDF5FilePusher::GetConnectedWriters() const {
    return 1;
}