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

#include "CBORStream2Serializer.h"
#include "tinycbor/cbor.h"
#include "CborErr.h"
#include "CborUtil.h"
#include "../compression/JFJochCompressor.h"
#include <nlohmann/json.hpp>

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const char* value) {
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encode_text_stringz(&encoder, value));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const std::string &value) {
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encode_text_stringz(&encoder, value.c_str()));
}

inline void CBOR_ENC_DATE(CborEncoder &encoder, const char* key, const std::string &value) {
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cbor_encode_tag(&encoder, CborDateTimeStringTag);
    cborErr(cbor_encode_text_stringz(&encoder, value.c_str()));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, float value) {
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encode_float(&encoder, value));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, bool value) {
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encode_boolean(&encoder, value));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, uint64_t value) {
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encode_uint(&encoder, value));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, int64_t value) {
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encode_int(&encoder, value));
}

template <class T>
void CBOR_ENC(CborEncoder &encoder, const char* key, const std::optional<T> &value) {
    if (value)
        CBOR_ENC(encoder, key, value.value());
}

void CBOR_ENC_COMPRESSED(CborEncoder &encoder,
                         const void *image, size_t image_size,
                         CompressionAlgorithm algorithm,
                         size_t elem_size) {
    if (algorithm == CompressionAlgorithm::NO_COMPRESSION)
        cborErr(cbor_encode_byte_string(&encoder, (uint8_t *) image, image_size));
    else {
        cbor_encode_tag(&encoder, TagDECTRISCompression);

        CborEncoder arrayEncoder;
        cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, 3));

        switch (algorithm) {
            case CompressionAlgorithm::BSHUF_LZ4:
                cborErr(cbor_encode_text_stringz(&arrayEncoder, "bslz4"));
                break;
            case CompressionAlgorithm::BSHUF_ZSTD:
            case CompressionAlgorithm::BSHUF_ZSTD_RLE:
                cborErr(cbor_encode_text_stringz(&arrayEncoder, "bszstd"));
                break;
            default:
                throw JFJochException(JFJochExceptionCategory::CBORError, "Unsupported compression algorithm");
        }
        cborErr(cbor_encode_uint(&arrayEncoder, elem_size));
        cborErr(cbor_encode_byte_string(&arrayEncoder, (uint8_t *) image, image_size));

        cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
    }
}

inline void CBOR_ENC_2D_TYPED_ARRAY(CborEncoder &encoder, const CompressedImage& image) {
    //if ((algorithm == CompressionAlgorithm::NO_COMPRESSION) && (xpixel * ypixel != image_size / elem_size))
    //    throw JFJochException(JFJochExceptionCategory::CBORError, "Mismatch in array size");

    CborEncoder arrayEncoder, arrayEncoder_2;
    cborErr(cbor_encode_text_stringz(&encoder, image.GetChannel().c_str()));

    cbor_encode_tag(&encoder, TagMultiDimArray);

    cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, 2));

    if (image.GetMode() == CompressedImageMode::RGB) {
        cborErr(cbor_encoder_create_array(&arrayEncoder, &arrayEncoder_2, 3));
        cborErr(cbor_encode_uint(&arrayEncoder_2, 3));
        cborErr(cbor_encode_uint(&arrayEncoder_2, image.GetHeight()));
        cborErr(cbor_encode_uint(&arrayEncoder_2, image.GetWidth()));
        cborErr(cbor_encoder_close_container(&arrayEncoder, &arrayEncoder_2));
    } else {
        cborErr(cbor_encoder_create_array(&arrayEncoder, &arrayEncoder_2, 2));
        cborErr(cbor_encode_uint(&arrayEncoder_2, image.GetHeight()));
        cborErr(cbor_encode_uint(&arrayEncoder_2, image.GetWidth()));
        cborErr(cbor_encoder_close_container(&arrayEncoder, &arrayEncoder_2));
    }

    CborTag typed_array_tag;

    switch (image.GetMode()) {
        case CompressedImageMode::RGB:
        case CompressedImageMode::Uint8:
            typed_array_tag = TagUnsignedInt8Bit;
            break;
        case CompressedImageMode::Uint16:
            typed_array_tag = TagUnsignedInt16BitLE;
            break;
        case CompressedImageMode::Uint32:
            typed_array_tag = TagUnsignedInt32BitLE;
            break;
        case CompressedImageMode::Int8:
            typed_array_tag = TagSignedInt8Bit;
            break;
        case CompressedImageMode::Int16:
            typed_array_tag = TagSignedInt16BitLE;
            break;
        case CompressedImageMode::Int32:
            typed_array_tag = TagSignedInt32BitLE;
            break;
        case CompressedImageMode::Float16:
            typed_array_tag = TagHalfLE;
            break;
        case CompressedImageMode::Float32:
            typed_array_tag = TagFloatLE;
            break;
        case CompressedImageMode::Float64:
            typed_array_tag = TagDoubleLE;
            break;
        default:
            throw JFJochException(JFJochExceptionCategory::CBORError, "Image mode not supported");
    }

    cbor_encode_tag(&arrayEncoder, typed_array_tag);

    CBOR_ENC_COMPRESSED(arrayEncoder, image.GetCompressed(),
                        image.GetCompressedSize(), image.GetCompressionAlgorithm(),
                        image.GetByteDepth());
    cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const std::vector<float>& v) {
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encode_tag(&encoder, TagFloatLE));
    cborErr(cbor_encode_byte_string(&encoder, (uint8_t *) v.data(), v.size() * sizeof(float)));
}


inline void CBOR_ENC(CborEncoder &encoder, const char* key, const std::vector<uint64_t>& v) {
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encode_tag(&encoder, TagUnsignedInt64BitLE));
    cborErr(cbor_encode_byte_string(&encoder, (uint8_t *) v.data(), v.size() * sizeof(uint64_t)));
}

inline void CBOR_ENC_RATIONAL(CborEncoder &encoder, const char* key, uint64_t numerator, uint64_t denominator) {
    CborEncoder arrayEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, 2));
    cborErr(cbor_encode_uint(&arrayEncoder, numerator));
    cborErr(cbor_encode_uint(&arrayEncoder, denominator));
    cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const std::map<std::string, float> &val) {
    CborEncoder mapEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, val.size()));
    for (auto &[map_key, map_val]: val)
        CBOR_ENC(mapEncoder, map_key.c_str(), map_val);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC_RAD_INT_RESULT(CborEncoder &encoder, const char* key,
                                    const std::map<std::string, std::vector<float>> &az_int_result) {
    CborEncoder mapEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, az_int_result.size()));
    for (auto &[map_key, map_val]: az_int_result)
        CBOR_ENC(mapEncoder, map_key.c_str(), map_val);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC_ADU_HIST(CborEncoder &encoder, const char* key,
                              const std::map<std::string, std::vector<uint64_t>> &adu_histogram) {
    CborEncoder mapEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, adu_histogram.size()));
    for (auto &[map_key, map_val]: adu_histogram)
        CBOR_ENC(mapEncoder, map_key.c_str(), map_val);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const SpotToSave& spot) {
    CborEncoder mapEncoder;
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, CborIndefiniteLength));
    CBOR_ENC(mapEncoder, "x", spot.x);
    CBOR_ENC(mapEncoder, "y", spot.y);
    CBOR_ENC(mapEncoder, "I", spot.intensity);
    CBOR_ENC(mapEncoder, "maxc", spot.maxc);
    CBOR_ENC(mapEncoder, "ice_ring", spot.ice_ring);
    CBOR_ENC(mapEncoder, "indexed", spot.indexed);
    if (spot.indexed) {
        CBOR_ENC(mapEncoder, "h", spot.h);
        CBOR_ENC(mapEncoder, "k", spot.k);
        CBOR_ENC(mapEncoder, "l", spot.l);
        CBOR_ENC(mapEncoder, "dist_ewald", spot.dist_ewald_sphere);
    }
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const XrayFluorescenceSpectrum& f) {
    if (f.empty())
        return;

    CborEncoder mapEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, CborIndefiniteLength));
    CBOR_ENC(mapEncoder, "data", f.GetData());
    CBOR_ENC(mapEncoder, "energy", f.GetEnergy_eV());
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const Reflection& r) {
    CborEncoder mapEncoder;
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, CborIndefiniteLength));
    CBOR_ENC(mapEncoder, "h", r.h);
    CBOR_ENC(mapEncoder, "k", r.k);
    CBOR_ENC(mapEncoder, "l", r.l);
    CBOR_ENC(mapEncoder, "x", r.predicted_x);
    CBOR_ENC(mapEncoder, "y", r.predicted_y);
    CBOR_ENC(mapEncoder, "d", r.d);
    CBOR_ENC(mapEncoder, "I", r.I);
    CBOR_ENC(mapEncoder, "bkg", r.bkg);
    CBOR_ENC(mapEncoder, "sigma", r.sigma);
    CBOR_ENC(mapEncoder, "image", r.image_number);
    CBOR_ENC(mapEncoder, "rp", r.dist_ewald);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const Coord& coord) {
    CborEncoder arrayEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, 3));
    cborErr(cbor_encode_float(&arrayEncoder, coord[0]));
    cborErr(cbor_encode_float(&arrayEncoder, coord[1]));
    cborErr(cbor_encode_float(&arrayEncoder, coord[2]));
    cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const std::vector<SpotToSave>& spots) {
    CborEncoder arrayEncoder, mapEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, spots.size()));

    for (auto spot : spots)
        CBOR_ENC(arrayEncoder, spot);

    cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const std::vector<Reflection>& refs) {
    CborEncoder arrayEncoder, mapEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, refs.size()));

    for (auto r: refs)
        CBOR_ENC(arrayEncoder, r);

    cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const CompressedImage& message) {
    CborEncoder mapEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 1));
    CBOR_ENC_2D_TYPED_ARRAY(mapEncoder, message);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC_AXIS(CborEncoder &encoder, const char* key, const float det_translation[3]) {
    CborEncoder arrayEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, 3));
    for (int i = 0; i < 3; i++)
        cborErr(cbor_encode_float(&arrayEncoder, det_translation[i]));
    cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
}

inline void CBOR_ENC_GONIOMETER(CborEncoder &encoder, const GoniometerAxis &g) {
    CborEncoder mapEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, g.GetName().c_str()));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, CborIndefiniteLength));
    CBOR_ENC(mapEncoder, "increment", g.GetIncrement_deg());
    CBOR_ENC(mapEncoder, "start", g.GetStart_deg());
    CBOR_ENC(mapEncoder, "axis", g.GetAxis());
    CBOR_ENC(mapEncoder, "screening_wedge", g.GetScreeningWedge());
    if (g.GetHelicalStep().has_value())
        CBOR_ENC(mapEncoder, "helical_step", g.GetHelicalStep().value());
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC_GRID_SCAN(CborEncoder &encoder, const char* key, const GridScanSettings &g) {
    CborEncoder mapEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 6));
    CBOR_ENC(mapEncoder, "n_fast", g.GetNFast());
    CBOR_ENC(mapEncoder, "n_slow", g.GetNSlow());
    CBOR_ENC(mapEncoder, "step_x_axis", g.GetGridStepX_um() * 1e-6f);
    CBOR_ENC(mapEncoder, "step_y_axis", g.GetGridStepY_um() * 1e-6f);
    CBOR_ENC(mapEncoder, "snake_scan", g.IsSnakeScan());
    CBOR_ENC(mapEncoder, "vertical_scan", g.IsVerticalScan());
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC_GONIOMETER_MAP(CborEncoder &encoder, const char* key, const StartMessage &msg) {
    CborEncoder mapEncoder;

    if (msg.goniometer) {
        cborErr(cbor_encode_text_stringz(&encoder, key));
        cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 1));
        CBOR_ENC_GONIOMETER(mapEncoder, msg.goniometer.value());
        cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
    }
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const std::vector<std::string> &v) {
    CborEncoder arrayEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, v.size()));
    for (const auto &i: v)
        cborErr(cbor_encode_text_stringz(&arrayEncoder, i.c_str()));
    cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const ROIMessage &val) {
    CborEncoder mapEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 6));
    CBOR_ENC(mapEncoder, "sum", val.sum);
    CBOR_ENC(mapEncoder, "sum_square", val.sum_square);
    CBOR_ENC(mapEncoder, "max_count", val.max_count);
    CBOR_ENC(mapEncoder, "pixels", val.pixels);
    CBOR_ENC(mapEncoder, "x_weighted_sum", val.x_weighted);
    CBOR_ENC(mapEncoder, "y_weighted_sum", val.y_weighted);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const std::map<std::string, ROIMessage> &map) {
    CborEncoder mapEncoder;

    if (map.empty())
        return;

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, map.size()));
    for (const auto &[x, y]: map)
        CBOR_ENC(mapEncoder, x.c_str(), y);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC_PIXEL_MASK(CborEncoder &encoder, const StartMessage &msg) {
    if (msg.pixel_mask.empty())
        return;

    CborEncoder mapEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, "pixel_mask"));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, msg.pixel_mask.size()));
    for (const auto &[key, value]: msg.pixel_mask) {
        if (value.size() != msg.image_size_x * msg.image_size_y)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Mismatch in size of pixel mask");

        JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
        auto mask_compressed = compressor.Compress(value);
        CompressedImage image(mask_compressed.data(), mask_compressed.size(),
                              msg.image_size_x, msg.image_size_y,
                              CompressedImageMode::Uint32,
                              CompressionAlgorithm::BSHUF_LZ4,
                              key);
        CBOR_ENC_2D_TYPED_ARRAY(mapEncoder, image);
    }
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC(CborEncoder &encoder, const char* key, const UnitCell &val) {
    CborEncoder mapEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 6));
    CBOR_ENC(mapEncoder, "a", val.a);
    CBOR_ENC(mapEncoder, "b", val.b);
    CBOR_ENC(mapEncoder, "c", val.c);
    CBOR_ENC(mapEncoder, "alpha", val.alpha);
    CBOR_ENC(mapEncoder, "beta", val.beta);
    CBOR_ENC(mapEncoder, "gamma", val.gamma);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

// Add encoder for LatticeMessage
inline void CBOR_ENC(CborEncoder &encoder, const char* key, const LatticeMessage &val) {
    CborEncoder mapEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 4));
    CBOR_ENC(mapEncoder, "centering", std::string(1, val.centering));
    CBOR_ENC(mapEncoder, "niggli_class", val.niggli_class);
    const char* cs_name = "triclinic";
    switch (val.crystal_system) {
        case gemmi::CrystalSystem::Triclinic:   cs_name = "triclinic";   break;
        case gemmi::CrystalSystem::Monoclinic:  cs_name = "monoclinic";  break;
        case gemmi::CrystalSystem::Orthorhombic:cs_name = "orthorhombic";break;
        case gemmi::CrystalSystem::Tetragonal:  cs_name = "tetragonal";  break;
        case gemmi::CrystalSystem::Trigonal:    cs_name = "trigonal";    break;
        case gemmi::CrystalSystem::Hexagonal:   cs_name = "hexagonal";   break;
        case gemmi::CrystalSystem::Cubic:       cs_name = "cubic";       break;
    }
    CBOR_ENC(mapEncoder, "system", std::string(cs_name));
    CBOR_ENC(mapEncoder, "primitive_lattice", val.primitive.GetVector());
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline nlohmann::json CBOR_ENC_ROI_CONFIG(const std::vector<ROIConfig> &roi) {
    nlohmann::json j;
    for (const auto &r: roi) {
        nlohmann::json jr;
        jr["name"] = r.name;
        switch (r.type) {
            case ROIConfig::ROIType::Box:
                jr["type"] = "box";
                jr["xmin"] = r.box.xmin;
                jr["xmax"] = r.box.xmax;
                jr["ymin"] = r.box.ymin;
                jr["ymax"] = r.box.ymax;
                break;
            case ROIConfig::ROIType::Circle:
                jr["type"] = "circle";
                jr["r"] = r.circle.r;
                jr["x"] = r.circle.x;
                jr["y"] = r.circle.y;
                break;
            case ROIConfig::ROIType::Azim:
                jr["type"] = "azim";
                jr["qmin"] = r.azim.qmin;
                jr["qmax"] = r.azim.qmax;
                break;
        }
        j.push_back(jr);
    }
    return j;
}

inline void CBOR_ENC_START_USER_DATA(CborEncoder& encoder, const char* key,
                                     const StartMessage& message) {
    nlohmann::json j;
    j["file_prefix"] = message.file_prefix;
    j["images_per_file"] = message.images_per_file;
    j["source_name"] = message.source_name;
    if (!message.source_type.empty())
        j["source_type"] = message.source_type;
    j["instrument_name"] = message.instrument_name;
    j["sample_name"] = message.sample_name;
    if (!message.user_data.empty())
        j["user"] = message.user_data;
    if (message.attenuator_transmission)
        j["attenuator_transmission"] = message.attenuator_transmission.value();
    if (message.total_flux)
        j["total_flux"] = message.total_flux.value();
    if (message.space_group_number)
        j["space_group_number"] = message.space_group_number.value();

    j["roi"] = CBOR_ENC_ROI_CONFIG(message.rois);

    j["gain_file_names"] = message.gain_file_names;
    if (message.write_master_file)
        j["write_master_file"] = message.write_master_file.value();
    if (message.data_reduction_factor_serialmx)
        j["data_reduction_factor_serialmx"] = message.data_reduction_factor_serialmx.value();
    j["experiment_group"] = message.experiment_group;
    j["jfjoch_release"] = message.jfjoch_release;
    if (message.socket_number)
        j["socket_number"] = message.socket_number.value();
    if (message.bit_depth_readout)
        j["bit_depth_readout"] = message.bit_depth_readout.value();
    if (!message.writer_notification_zmq_addr.empty())
        j["writer_notification_zmq_addr"] = message.writer_notification_zmq_addr;
    if (message.summation_mode.has_value())
        j["summation_mode"] = message.summation_mode.value();
    if (message.overwrite.has_value())
        j["overwrite"] = message.overwrite.value();
    if (message.xfel_pulse_id.has_value())
        j["xfel_pulse_id"] = message.xfel_pulse_id.value();
    if (message.ring_current_mA.has_value())
        j["ring_current_mA"] = message.ring_current_mA.value();
    if (message.sample_temperature_K.has_value())
        j["sample_temperature_K"] = message.sample_temperature_K.value();
    if (message.file_format.has_value())
        j["file_format"] = static_cast<int>(message.file_format.value());

    if (message.images_per_trigger.has_value())
        j["images_per_trigger"] = message.images_per_trigger.value();

    if (message.poni_rot1.has_value())
        j["poni_rot1"] = message.poni_rot1.value();
    if (message.poni_rot2.has_value())
        j["poni_rot2"] = message.poni_rot2.value();
    if (message.poni_rot3.has_value())
        j["poni_rot3"] = message.poni_rot3.value();

    if (message.detect_ice_rings.has_value())
        j["detect_ice_rings"] = message.detect_ice_rings.value();

    switch(message.indexing_algorithm) {
        case IndexingAlgorithmEnum::FFBIDX:
            j["indexing_algorithm"] = "ffbidx";
            break;
        case IndexingAlgorithmEnum::FFT:
            j["indexing_algorithm"] = "fft";
            break;
        case IndexingAlgorithmEnum::FFTW:
            j["indexing_algorithm"] = "fftw";
            break;
        default:
            j["indexing_algorithm"] = "none";
            break;
    }

    switch (message.geom_refinement_algorithm) {
        case GeomRefinementAlgorithmEnum::BeamCenter:
            j["geom_refinement_algorithm"] = "beam_center";
            break;
        default:
            j["geom_refinement_algorithm"] = "none";
            break;
    }

    auto str = j.dump();

    CBOR_ENC(encoder, key, str);
}

CBORStream2Serializer::CBORStream2Serializer(uint8_t *in_buffer, size_t buffer_size) :
buffer(in_buffer), max_buffer_size(buffer_size), curr_size(0) {}

size_t CBORStream2Serializer::GetBufferSize() const {
    return curr_size;
}

void CBORStream2Serializer::SerializeSequenceStart(const StartMessage& message) {
    CborEncoder encoder, mapEncoder;

    cbor_encoder_init(&encoder, buffer, max_buffer_size, 0);
    cborErr(cbor_encode_tag(&encoder, CborSignatureTag ));

    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, CborIndefiniteLength));

    CBOR_ENC(mapEncoder, "type", "start");
    CBOR_ENC(mapEncoder, "magic_number", user_data_magic_number);

    CBOR_ENC(mapEncoder, "detector_distance", message.detector_distance);
    CBOR_ENC_AXIS(mapEncoder, "detector_translation", message.detector_translation);
    CBOR_ENC(mapEncoder, "beam_center_x", message.beam_center_x);
    CBOR_ENC(mapEncoder, "beam_center_y", message.beam_center_y);
    CBOR_ENC(mapEncoder, "countrate_correction_enabled", message.countrate_correction_enabled);
    CBOR_ENC(mapEncoder, "flatfield_enabled", message.flatfield_enabled);
    CBOR_ENC(mapEncoder, "number_of_images", message.number_of_images);
    CBOR_ENC(mapEncoder, "image_size_x", message.image_size_x);
    CBOR_ENC(mapEncoder, "image_size_y", message.image_size_y);

    CBOR_ENC(mapEncoder, "incident_energy", message.incident_energy);
    CBOR_ENC(mapEncoder, "incident_wavelength", message.incident_wavelength);

    CBOR_ENC(mapEncoder, "frame_time", message.frame_time);
    CBOR_ENC(mapEncoder, "count_time", message.count_time);

    CBOR_ENC(mapEncoder, "saturation_value", message.saturation_value);
    CBOR_ENC(mapEncoder, "error_value", message.error_value);
    CBOR_ENC(mapEncoder, "pixel_size_x", message.pixel_size_x);
    CBOR_ENC(mapEncoder, "pixel_size_y", message.pixel_size_y);
    CBOR_ENC(mapEncoder, "sensor_thickness", message.sensor_thickness);
    CBOR_ENC(mapEncoder, "sensor_material", message.sensor_material);
    CBOR_ENC_DATE(mapEncoder, "arm_date", message.arm_date);
    CBOR_ENC(mapEncoder, "pixel_mask_enabled", message.pixel_mask_enabled);
    CBOR_ENC(mapEncoder, "detector_description", message.detector_description);
    CBOR_ENC(mapEncoder, "detector_serial_number", message.detector_serial_number);
    CBOR_ENC(mapEncoder, "series_unique_id", message.run_name);
    CBOR_ENC(mapEncoder, "series_id", message.run_number);
    CBOR_ENC(mapEncoder, "fluorescence", message.fluorescence_spectrum);

    if (message.goniometer)
        CBOR_ENC_GONIOMETER_MAP(mapEncoder, "goniometer", message);
    else if (message.grid_scan)
        CBOR_ENC_GRID_SCAN(mapEncoder, "grid_scan", message.grid_scan.value());

    CBOR_ENC(mapEncoder, "jungfrau_conversion_enabled", message.jungfrau_conversion_enabled);
    CBOR_ENC(mapEncoder, "jungfrau_conversion_factor", message.jungfrau_conversion_factor);
    CBOR_ENC(mapEncoder, "geometry_transformation_enabled", message.geometry_transformation_enabled);

    CBOR_ENC_PIXEL_MASK(mapEncoder, message);

    CBOR_ENC(mapEncoder, "channels", message.channels);
    CBOR_ENC(mapEncoder, "max_spot_count", message.max_spot_count);

    CBOR_ENC(mapEncoder, "storage_cell_number", message.storage_cell_number);
    CBOR_ENC_RATIONAL(mapEncoder, "storage_cell_delay", message.storage_cell_delay_ns, 1000*1000*1000UL);
    CBOR_ENC(mapEncoder, "threshold_energy", message.threshold_energy);

    switch (message.bit_depth_image) {
        case 8:
            CBOR_ENC(mapEncoder, "image_dtype", message.pixel_signed ? "int8" : "uint8");
            break;
        case 16:
            CBOR_ENC(mapEncoder, "image_dtype", message.pixel_signed ? "int16" : "uint16");
            break;
        case 32:
            CBOR_ENC(mapEncoder, "image_dtype", message.pixel_signed ? "int32" : "uint32");
            break;
    }

    CBOR_ENC(mapEncoder, "unit_cell", message.unit_cell);
    CBOR_ENC(mapEncoder, "az_int_q_bin_count", message.az_int_q_bin_count);
    CBOR_ENC(mapEncoder, "az_int_phi_bin_count", message.az_int_phi_bin_count);
    CBOR_ENC(mapEncoder, "az_int_bin_to_q", message.az_int_bin_to_q);
    CBOR_ENC(mapEncoder, "az_int_bin_to_two_theta", message.az_int_bin_to_two_theta);
    if (!message.az_int_bin_to_phi.empty())
        CBOR_ENC(mapEncoder, "az_int_bin_to_phi", message.az_int_bin_to_phi);
    CBOR_ENC(mapEncoder, "summation", message.summation);

    CBOR_ENC_START_USER_DATA(mapEncoder, "user_data", message);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
    curr_size = cbor_encoder_get_buffer_size(&encoder, buffer);
}

void CBORStream2Serializer::SerializeSequenceEnd(const EndMessage& message) {
    CborEncoder encoder, mapEncoder;
    cbor_encoder_init(&encoder, buffer, max_buffer_size, 0);
    cborErr(cbor_encode_tag(&encoder, CborSignatureTag));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, CborIndefiniteLength));

    CBOR_ENC(mapEncoder, "type", "end");
    CBOR_ENC(mapEncoder, "magic_number", user_data_magic_number);

    CBOR_ENC(mapEncoder, "series_unique_id", message.run_name);
    CBOR_ENC(mapEncoder, "series_id", message.run_number);
    CBOR_ENC(mapEncoder, "end_date", message.end_date);

    CBOR_ENC(mapEncoder, "max_image_number", message.max_image_number);
    CBOR_ENC(mapEncoder, "images_collected", message.images_collected_count);
    CBOR_ENC(mapEncoder, "images_sent_to_write", message.images_sent_to_write_count);
    CBOR_ENC(mapEncoder, "data_collection_efficiency", message.efficiency);
    CBOR_ENC_RAD_INT_RESULT(mapEncoder, "az_int_result", message.az_int_result);
    CBOR_ENC_ADU_HIST(mapEncoder, "adu_histogram", message.adu_histogram);
    CBOR_ENC(mapEncoder, "adu_histogram_bin_width", message.adu_histogram_bin_width);
    CBOR_ENC(mapEncoder, "max_receiver_delay", message.max_receiver_delay);
    CBOR_ENC(mapEncoder, "indexing_rate", message.indexing_rate);
    CBOR_ENC(mapEncoder, "bkg_estimate", message.bkg_estimate);

    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));

    curr_size = cbor_encoder_get_buffer_size(&encoder, buffer);
}

void CBORStream2Serializer::SerializeImageInternal(CborEncoder &mapEncoder, const DataMessage &message, bool metadata_only) {
    CBOR_ENC(mapEncoder, "image_id", message.number);

    CBOR_ENC(mapEncoder, "original_image_id", message.original_number);

    CBOR_ENC_RATIONAL(mapEncoder, "real_time", message.exptime, message.exptime_base);
    CBOR_ENC_RATIONAL(mapEncoder, "start_time", message.timestamp, message.timestamp_base);
    CBOR_ENC_RATIONAL(mapEncoder, "end_time", message.timestamp + message.exptime, message.timestamp_base);

    CBOR_ENC(mapEncoder, "spot_count", message.spot_count);
    CBOR_ENC(mapEncoder, "spot_count_ice_rings", message.spot_count_ice_rings);
    CBOR_ENC(mapEncoder, "spot_count_low_res", message.spot_count_low_res);
    CBOR_ENC(mapEncoder, "spot_count_indexed", message.spot_count_indexed);
    CBOR_ENC(mapEncoder, "az_int_profile", message.az_int_profile);
    CBOR_ENC(mapEncoder, "indexing_result", message.indexing_result);
    if (message.indexing_lattice)
        CBOR_ENC(mapEncoder, "indexing_lattice", message.indexing_lattice->GetVector());
    CBOR_ENC(mapEncoder, "profile_radius", message.profile_radius);
    CBOR_ENC(mapEncoder, "b_factor", message.b_factor);
    CBOR_ENC(mapEncoder, "indexing_time", message.indexing_time_s);
    CBOR_ENC(mapEncoder, "processing_time", message.processing_time_s);
    CBOR_ENC(mapEncoder, "indexing_unit_cell", message.indexing_unit_cell);
    CBOR_ENC(mapEncoder, "xfel_pulse_id", message.xfel_pulse_id);
    CBOR_ENC(mapEncoder, "xfel_event_code", message.xfel_event_code);
    if (message.lattice_type)
        CBOR_ENC(mapEncoder, "lattice_type", message.lattice_type.value());
    CBOR_ENC(mapEncoder, "jf_info", message.jf_info);
    CBOR_ENC(mapEncoder, "receiver_aq_dev_delay", message.receiver_aq_dev_delay);
    CBOR_ENC(mapEncoder, "receiver_free_send_buf", message.receiver_free_send_buf);
    CBOR_ENC(mapEncoder, "storage_cell", message.storage_cell);
    CBOR_ENC(mapEncoder, "saturated_pixel_count", message.saturated_pixel_count);
    CBOR_ENC(mapEncoder, "pixel_sum", message.pixel_sum);
    CBOR_ENC(mapEncoder, "error_pixel_count", message.error_pixel_count);
    CBOR_ENC(mapEncoder, "strong_pixel_count", message.strong_pixel_count);
    CBOR_ENC(mapEncoder, "min_viable_pixel_value", message.min_viable_pixel_value);
    CBOR_ENC(mapEncoder, "max_viable_pixel_value", message.max_viable_pixel_value);
    CBOR_ENC(mapEncoder, "resolution_estimate", message.resolution_estimate);
    CBOR_ENC(mapEncoder, "data_collection_efficiency", message.image_collection_efficiency);
    CBOR_ENC(mapEncoder, "packets_expected", message.packets_expected);
    CBOR_ENC(mapEncoder, "packets_received", message.packets_received);
    CBOR_ENC(mapEncoder, "bkg_estimate", message.bkg_estimate);
    CBOR_ENC(mapEncoder, "adu_histogram", message.adu_histogram);
    CBOR_ENC(mapEncoder, "roi_integrals", message.roi);
    CBOR_ENC(mapEncoder, "beam_corr_x", message.beam_corr_x);
    CBOR_ENC(mapEncoder, "beam_corr_y", message.beam_corr_y);
    CBOR_ENC(mapEncoder, "user_data", message.user_data.dump());

    if (!metadata_only) {
        CBOR_ENC(mapEncoder, "spots", message.spots);
        CBOR_ENC(mapEncoder, "reflections", message.reflections);
        CBOR_ENC(mapEncoder, "data", message.image);
    }
}

void CBORStream2Serializer::SerializeImage(const DataMessage& message) {
    CborEncoder encoder, mapEncoder;
    cbor_encoder_init(&encoder, buffer, max_buffer_size, 0);

    cborErr(cbor_encode_tag(&encoder, CborSignatureTag ));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, CborIndefiniteLength));

    CBOR_ENC(mapEncoder, "type", "image");
    CBOR_ENC(mapEncoder, "magic_number", user_data_magic_number);
    CBOR_ENC(mapEncoder, "series_unique_id", message.run_name);
    CBOR_ENC(mapEncoder, "series_id", message.run_number);

    SerializeImageInternal(mapEncoder, message, false);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));

    curr_size = cbor_encoder_get_buffer_size(&encoder, buffer);
}

void CBORStream2Serializer::SerializeMetadata(const MetadataMessage &messages) {
    if (messages.images.empty())
        throw JFJochException(JFJochExceptionCategory::CBORError,
            "Cannot serialize empty metadata packet");
    CborEncoder encoder, mapEncoder, arrayEncoder;
    cbor_encoder_init(&encoder, buffer, max_buffer_size, 0);

    cborErr(cbor_encode_tag(&encoder, CborSignatureTag ));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, CborIndefiniteLength));
    CBOR_ENC(mapEncoder, "type", "metadata");
    CBOR_ENC(mapEncoder, "magic_number", user_data_magic_number);
    CBOR_ENC(mapEncoder, "series_unique_id", messages.run_name);
    CBOR_ENC(mapEncoder, "series_id", messages.run_number);
    cborErr(cbor_encode_text_stringz(&mapEncoder, "images"));
    cborErr(cbor_encoder_create_array(&mapEncoder, &arrayEncoder, messages.images.size()));

    for (const auto &image: messages.images) {
        CborEncoder localEncoder;

        cborErr(cbor_encoder_create_map(&arrayEncoder, &localEncoder, CborIndefiniteLength));

        SerializeImageInternal(localEncoder, image, true);
        cborErr(cbor_encoder_close_container(&arrayEncoder, &localEncoder));    
    }
    
    cborErr(cbor_encoder_close_container(&mapEncoder, &arrayEncoder));
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
    curr_size = cbor_encoder_get_buffer_size(&encoder, buffer);
}

void CBORStream2Serializer::SerializeCalibration(const CompressedImage &image) {
    CborEncoder encoder, mapEncoder;
    cbor_encoder_init(&encoder, buffer, max_buffer_size, 0);

    cborErr(cbor_encode_tag(&encoder, CborSignatureTag ));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, CborIndefiniteLength));
    CBOR_ENC(mapEncoder, "type", "calibration");
    CBOR_ENC(mapEncoder, "magic_number", user_data_magic_number);
    CBOR_ENC(mapEncoder, "data", image);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));

    curr_size = cbor_encoder_get_buffer_size(&encoder, buffer);
}

size_t CBORStream2Serializer::GetImageAppendOffset() const {
    return curr_size + sizeof(size_t) - 1;
}

void CBORStream2Serializer::AppendImage(size_t image_size) {
    if (curr_size + image_size + sizeof(size_t) + 1 >= max_buffer_size)
        throw JFJochException(JFJochExceptionCategory::CBORError, "No space to extend the image");

    buffer[curr_size - 2] = 0x40 + 27;
    curr_size--;
#ifdef LITTLE_ENDIAN
    size_t image_size_be = __builtin_bswap64(image_size);
#else
    size_t image_size_be = image_size;
#endif
    memcpy(buffer + curr_size, &image_size_be, sizeof(size_t));
    curr_size += sizeof(size_t);
    curr_size += image_size + 0;
    buffer[curr_size] = 0xFF;
    curr_size++;
}