// Copyright (2019-2023) Paul Scherrer Institute

#include "ZMQBsreadImagePusher.h"
#include <nlohmann/json.hpp>
#include <openssl/md5.h>

ZMQBsreadImagePusher::ZMQBsreadImagePusher(const std::string &addr, int32_t send_buffer_high_watermark,
                                           int32_t send_buffer_size) : socket(context, ZMQSocketType::Push) {
    if (send_buffer_size > 0)
        socket.SendBufferSize(send_buffer_size);
    if (send_buffer_high_watermark > 0)
        socket.SendWaterMark(send_buffer_high_watermark);
    socket.Bind(addr);
}

void ZMQBsreadImagePusher::StartDataCollection(const StartMessage &message) {
    // Do nothing
}

bool ZMQBsreadImagePusher::SendImage(const DataMessage &message) {
    std::unique_lock<std::mutex> ul(m);

    timespec ts{};
    clock_gettime(CLOCK_REALTIME, &ts);
    auto data_header = DataHeader(message);
    auto main_header = MainHeader(message, data_header, ts);
    socket.Send(main_header, true, true);
    socket.Send(data_header, true, true);
    socket.Send(message.image.data, message.image.size, true, true);
    socket.Send(&ts, sizeof(ts), true, false);
    return true; // Blocking! - it would be unsafe to use non-blocking send with multipart
}

bool ZMQBsreadImagePusher::EndDataCollection(const EndMessage &message) {
    return true;
}

std::string ZMQBsreadImagePusher::MainHeader(const DataMessage &message, const std::string &data_header, timespec &ts) {
    nlohmann::json j;
    j["htype"] = "bsr_m-1.1";
    j["pulse_id"] = message.xfel_bunch_id;
    j["hash"] = MD5Hash(data_header);
    j["global_timestamp"]["sec"] = ts.tv_sec;
    j["global_timestamp"]["ns"] = ts.tv_nsec;
    return j.dump();
}

std::string ZMQBsreadImagePusher::MD5Hash(const std::string &s) {
    std::string ret;
    uint8_t hash[MD5_DIGEST_LENGTH];
    MD5(reinterpret_cast<const unsigned char *>(s.data()), s.length(), hash);
    for (unsigned char i : hash) {
        char buf[4];
        sprintf(buf, "%02x", i);
        ret.append(buf);
    }
    return ret;
}

std::string ZMQBsreadImagePusher::DataHeader(const DataMessage &message) {
    nlohmann::json data;
    data["htype"] = "bsr_d-1.1";

    nlohmann::json image_channel;
    image_channel["name"] = "JFJOCH:IMAGE";
    if (message.image.pixel_depth_bytes == 2)
        image_channel["type"] = (message.image.pixel_is_signed ? "int16" : "uint16");
    else
        image_channel["type"] = (message.image.pixel_is_signed ? "int32" : "uint32");
    image_channel["shape"] = {message.image.ypixel, message.image.xpixel};
    image_channel["encoding"] = "little";
    if (message.image.algorithm == CompressionAlgorithm::BSHUF_LZ4)
        image_channel["compression"] = "bitshuffle_lz4";
    else if (message.image.algorithm == CompressionAlgorithm::NO_COMPRESSION)
        image_channel["compression"] = "none";
    else
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Compression not allowed for BSRead data");
    image_channel["modulo"] = 1;
    image_channel["offset"] = 0;
    data["channels"].push_back(image_channel);
    return data.dump();
}