// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include "JFJochFrameSerializer.h"
#include "tinycbor/src/cbor.h"
#include "CborErr.h"
#include "CborUtil.h"

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));
    if (value.empty())
        cborErr(cbor_encode_text_stringz(&encoder, "."));
    else
        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);
    if (value.empty())
        cborErr(cbor_encode_text_stringz(&encoder, "."));
    else
        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));
}

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_MULTIDIM_TYPED_ARRAY(CborEncoder &encoder, const char* key,
                                              const void *image, size_t image_size,
                                              size_t xpixel, size_t ypixel, CompressionAlgorithm algorithm,
                                              size_t elem_size, bool elem_sign) {
    //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, key));

    cbor_encode_tag(&encoder, TagMultiDimArray);

    cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, 2));
    cborErr(cbor_encoder_create_array(&arrayEncoder, &arrayEncoder_2, 2));
    cborErr(cbor_encode_uint(&arrayEncoder_2, ypixel));
    cborErr(cbor_encode_uint(&arrayEncoder_2, xpixel));
    cborErr(cbor_encoder_close_container(&arrayEncoder, &arrayEncoder_2));

    CborTag typed_array_tag;

    if (elem_sign) {
        if (elem_size == 4)
            typed_array_tag = TagSignedInt32BitLE;
        else if (elem_size == 2)
            typed_array_tag = TagSignedInt16BitLE;
        else if (elem_size == 1)
            typed_array_tag = TagSignedInt8Bit;
        else
            throw JFJochException(JFJochExceptionCategory::CBORError, "Array size not supported");
    } else {
        if (elem_size == 4)
            typed_array_tag = TagUnsignedInt32BitLE;
        else if (elem_size == 2)
            typed_array_tag = TagUnsignedInt16BitLE;
        else if (elem_size == 1)
            typed_array_tag = TagUnsignedInt8Bit;
        else
            throw JFJochException(JFJochExceptionCategory::CBORError, "Array size not supported");
    }

    cbor_encode_tag(&arrayEncoder, typed_array_tag);

    CBOR_ENC_COMPRESSED(arrayEncoder, image, image_size, algorithm, elem_size);
    cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
}

inline void CBOR_ENC_PIXEL_MASK(CborEncoder &encoder, const char* key,
                                const std::map<std::string, std::vector<uint32_t>> &pixel_mask,
                                size_t xpixel, size_t ypixel) {
    CborEncoder mapEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, pixel_mask.size()));
    for (auto &[pixel_mask_key, pixel_mask_array]: pixel_mask)
        CBOR_ENC_MULTIDIM_TYPED_ARRAY(mapEncoder,
                                          pixel_mask_key.c_str(),
                                          pixel_mask_array.data(),
                                          pixel_mask_array.size() * sizeof(uint32_t),
                                          xpixel, ypixel,
                                          CompressionAlgorithm::NO_COMPRESSION,
                                          sizeof(uint32_t),
                                          false);

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

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

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encode_tag(&encoder, TagFloatLE));
    cborErr(cbor_encoder_create_array(&encoder, &arrayEncoder, v.size()));

    for (const auto &i : v)
        cborErr(cbor_encode_float(&arrayEncoder, i));

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

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_RAD_INT_RESULT(CborEncoder &encoder, const char* key,
                                    const std::map<std::string, std::vector<float>> &rad_int_result) {
    CborEncoder mapEncoder;
    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, rad_int_result.size()));
    for (auto &[map_key, map_val]: rad_int_result)
        CBOR_ENC(mapEncoder, map_key.c_str(), map_val);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

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

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

    for (auto spot : spots) {
        cborErr(cbor_encode_float(&arrayEncoder, spot.x));
        cborErr(cbor_encode_float(&arrayEncoder, spot.y));
        cborErr(cbor_encode_float(&arrayEncoder, spot.intensity));
        cborErr(cbor_encode_boolean(&arrayEncoder, spot.indexed));
    }
    cborErr(cbor_encoder_close_container(&encoder, &arrayEncoder));
}

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

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 1));
    CBOR_ENC_MULTIDIM_TYPED_ARRAY(mapEncoder,
                                  message.channel.c_str(),
                                  (uint8_t *) message.data,
                                  message.size,
                                  message.xpixel,
                                  message.ypixel,
                                  message.algorithm,
                                  message.pixel_depth_bytes,
                                  message.pixel_is_signed);

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

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

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 2));
    cborErr(cbor_encode_text_stringz(&mapEncoder, "increment"));
    cborErr(cbor_encode_float(&mapEncoder, g.increment));
    cborErr(cbor_encode_text_stringz(&mapEncoder, "start"));
    cborErr(cbor_encode_float(&mapEncoder, g.start));
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}


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

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, g.size()));
    for (const auto& [x, y]: g)
        CBOR_ENC_GONIOMETER(mapEncoder, x.c_str(), y);

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

inline void CBOR_ENC_DET_TRANSLATION(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_CHANNELS(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_UNIT_CELL(CborEncoder &encoder, const char* key, const float val[6]) {
    CborEncoder mapEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, key));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 6));
    CBOR_ENC(mapEncoder, "a", val[0]);
    CBOR_ENC(mapEncoder, "b", val[1]);
    CBOR_ENC(mapEncoder, "c", val[2]);
    CBOR_ENC(mapEncoder, "alpha", val[3]);
    CBOR_ENC(mapEncoder, "beta", val[4]);
    CBOR_ENC(mapEncoder, "gamma", val[5]);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

inline void CBOR_ENC_USER_DATA(CborEncoder &encoder, const StartMessage& message) {
    CborEncoder mapEncoder;

    cborErr(cbor_encode_text_stringz(&encoder, "user_data"));
    cborErr(cbor_encoder_create_map(&encoder, &mapEncoder, 20));

    CBOR_ENC(mapEncoder, "file_prefix", message.file_prefix);
    CBOR_ENC(mapEncoder, "sample_name", message.sample_name);
    CBOR_ENC(mapEncoder, "space_group_number", message.space_group_number);
    CBOR_ENC(mapEncoder, "max_spot_count", message.max_spot_count);
    CBOR_ENC(mapEncoder, "data_file_count", message.data_file_count);
    CBOR_ENC(mapEncoder, "storage_cell_number", message.storage_cell_number);
    CBOR_ENC(mapEncoder, "pixel_bit_depth", message.pixel_bit_depth);
    CBOR_ENC(mapEncoder, "pixel_signed", message.pixel_signed);
    CBOR_ENC(mapEncoder, "min_value", message.min_value);
    switch (message.compression_algorithm) {
        case CompressionAlgorithm::BSHUF_LZ4:
            CBOR_ENC(mapEncoder, "compression_algorithm", "bslz4");
            break;
        case CompressionAlgorithm::BSHUF_ZSTD:
        case CompressionAlgorithm::BSHUF_ZSTD_RLE:
            CBOR_ENC(mapEncoder, "compression_algorithm", "bszstd");
            break;
        default:
            CBOR_ENC(mapEncoder, "compression_algorithm", "none");
            break;
    }
    CBOR_ENC(mapEncoder, "compression_block_size", message.compression_block_size);

    CBOR_ENC_UNIT_CELL(mapEncoder, "unit_cell", message.unit_cell);

    CBOR_ENC(mapEncoder, "source_name", message.source_name);
    CBOR_ENC(mapEncoder, "source_name_short", message.source_name_short);
    CBOR_ENC(mapEncoder, "instrument_name", message.instrument_name);
    CBOR_ENC(mapEncoder, "instrument_name_short", message.instrument_name_short);
    CBOR_ENC(mapEncoder, "rad_int_bin_number", message.rad_int_bin_number);
    CBOR_ENC(mapEncoder, "rad_int_bin_to_q", message.rad_int_bin_to_q);
    CBOR_ENC(mapEncoder, "rad_int_solid_angle_corr", message.rad_int_solid_angle_corr);
    CBOR_ENC(mapEncoder, "summation", message.summation);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
}

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


size_t JFJochFrameSerializer::GetImageAppendOffset() {
    return  curr_size + sizeof(size_t);
}

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

size_t JFJochFrameSerializer::GetRemainingBuffer() const {
    return max_buffer_size - curr_size;
}

void JFJochFrameSerializer::AppendImage(size_t image_size) {
    if (curr_size + image_size + sizeof(size_t) >= max_buffer_size)
        throw JFJochException(JFJochExceptionCategory::CBORError, "No space to extend the image");
    curr_size--;
    buffer[curr_size] = 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;
}


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

    cbor_encoder_init(&encoder, buffer, max_buffer_size, 0);
    cborErr(cbor_encode_tag(&encoder, CborSignatureTag ));
    size_t elements = 27;

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

    CBOR_ENC(mapEncoder, "type", "start");
    CBOR_ENC(mapEncoder, "detector_distance", message.detector_distance);
    CBOR_ENC(mapEncoder, "beam_center_x", message.beam_center_x);
    CBOR_ENC(mapEncoder, "beam_center_y", message.beam_center_y);

    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, "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.series_unique_id);
    CBOR_ENC(mapEncoder, "series_id", message.series_id);

    CBOR_ENC_DET_TRANSLATION(mapEncoder, "detector_translation", message.detector_translation);
    CBOR_ENC_GONIOMETER_MAP(mapEncoder, "goniometer", message.goniometer);

    CBOR_ENC_USER_DATA(mapEncoder, message);
    CBOR_ENC_PIXEL_MASK(mapEncoder, "pixel_mask", message.pixel_mask, message.image_size_x, message.image_size_y);
    CBOR_ENC_CHANNELS(mapEncoder, "channels", message.channels);

    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));
    curr_size = cbor_encoder_get_buffer_size(&encoder, buffer);
}

void JFJochFrameSerializer::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, 9));

    CBOR_ENC(mapEncoder, "type", "end");
    CBOR_ENC(mapEncoder, "series_unique_id", message.series_unique_id);
    CBOR_ENC(mapEncoder, "series_id", message.series_id);
    CBOR_ENC(mapEncoder, "number_of_images", message.number_of_images);
    CBOR_ENC(mapEncoder, "max_receiver_delay", message.max_receiver_delay);
    CBOR_ENC(mapEncoder, "receiver_efficiency", message.efficiency);
    CBOR_ENC(mapEncoder, "write_master_file", message.write_master_file);
    CBOR_ENC(mapEncoder, "end_date", message.end_date);
    CBOR_ENC_RAD_INT_RESULT(mapEncoder, "rad_int_result", message.rad_int_result);
    cborErr(cbor_encoder_close_container(&encoder, &mapEncoder));

    curr_size = cbor_encoder_get_buffer_size(&encoder, buffer);
}

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

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

    CBOR_ENC(mapEncoder, "type", "image");
    CBOR_ENC(mapEncoder, "series_unique_id", message.series_unique_id);
    CBOR_ENC(mapEncoder, "series_id", message.series_id);
    CBOR_ENC(mapEncoder, "image_id", message.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);

    cborErr(cbor_encode_text_stringz(&mapEncoder, "user_data"));
    cborErr(cbor_encoder_create_map(&mapEncoder, &userDataMapEncoder, 9));

    CBOR_ENC(userDataMapEncoder, "spots", message.spots);
    CBOR_ENC(userDataMapEncoder, "rad_int_profile", message.rad_int_profile);
    CBOR_ENC(userDataMapEncoder, "indexing_result", message.indexing_result);
    CBOR_ENC(userDataMapEncoder, "indexing_lattice", message.indexing_lattice);
    CBOR_ENC(userDataMapEncoder, "bunch_id", message.bunch_id);
    CBOR_ENC(userDataMapEncoder, "jf_info", (uint64_t) message.jf_info);
    CBOR_ENC(userDataMapEncoder, "receiver_available_send_buffers", message.receiver_available_send_buffers);
    CBOR_ENC(userDataMapEncoder, "receiver_aq_dev_delay", message.receiver_aq_dev_delay);
    CBOR_ENC(userDataMapEncoder, "storage_cell", (uint64_t) message.storage_cell);
    cborErr(cbor_encoder_close_container(&mapEncoder, &userDataMapEncoder));

    CBOR_ENC(mapEncoder, "data", message.image);

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

    curr_size = cbor_encoder_get_buffer_size(&encoder, buffer);
}