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

#include "CBFWriter.h"

#include <utility>

#include "../include/spdlog/fmt/fmt.h"
#include "../compression/JFJochDecompress.h"

template <class T>
void write(uint8_t* arr, T val, size_t &offset) {
    auto ptr = reinterpret_cast<T *>(arr + offset);
    *ptr = val;
    offset += sizeof(T);
}

template <class T>
void CBFWriter::EncodeCBF(const T *image, size_t num_elem, std::vector<uint8_t> &output) {
    size_t output_loc = 0;
    int64_t cval = 0;

    for (int i = 0; i < num_elem; i++) {
        if (output_loc > output.size() - 6)
            output.resize(output.size() * 2);

        int64_t nval;
        if (image[i] > saturation_value)
            nval = saturation_value;
        else if ((image[i] < 0) || (pixel_mask[i] != 0))
            // For Bitshuffle compression INTx_MIN is the best value for special things, at it is one bit different from 0
            // For CBF compression -1 is much better, as it reduces differences in counts
            nval = -1;
        else
            nval = image[i];

        int64_t diff = nval - cval;
        if (abs(diff) <= 127)
            write(output.data(), static_cast<int8_t>(diff), output_loc);
        else {
            write(output.data(), static_cast<uint8_t>(0x80), output_loc);
            if (abs(diff) <= 32767)
                write(output.data(), static_cast<int16_t>(diff), output_loc);
            else {
                write(output.data(), static_cast<uint16_t>(0x8000), output_loc);
                write(output.data(), static_cast<int32_t>(diff), output_loc);
            }
        }
        cval = nval;      // current value
    }
    output.resize(output_loc);
}

CBFWriter::CBFWriter(StartMessage in_start_msg)
        : start_msg(std::move(in_start_msg)),
          saturation_value(std::min<int64_t>(start_msg.saturation_value, (1L<<24))),
          pixel_mask(start_msg.image_size_x * start_msg.image_size_y, 0) {

    header += fmt::format("\r\n# Detector: {:s}, S/N {:s}\r\n",
                          start_msg.detector_description,
                          start_msg.detector_serial_number);

    header += fmt::format("# Silicon sensor, thickness {:.6f} m\r\n", start_msg.sensor_thickness);
    header += fmt::format("# Exposure_time {:.6f} s\r\n", start_msg.count_time);
    header += fmt::format("# Exposure_period {:.6f} s\r\n", start_msg.frame_time);
    header += fmt::format("# Count_cutoff {:d} counts\r\n", saturation_value);
    header += fmt::format("# Wavelength {:.3f} A\r\n", start_msg.incident_wavelength);
    header += fmt::format("# Detector_distance {:.4f} m\r\n", start_msg.detector_distance);
    header += fmt::format("# Beam_xy ({:.1f}, {:.1f}) pixels\r\n", start_msg.beam_center_x, start_msg.beam_center_y);

    if (start_msg.goniometer.has_value())
        header += fmt::format("# Start_angle {:f} deg.\r\n# Angle_increment {:f} deg.\r\n",
                              start_msg.goniometer->GetStart_deg(),
                              start_msg.goniometer->GetIncrement_deg());

    if ((start_msg.pixel_mask.size() == 1) && (start_msg.pixel_mask.begin()->second.size() == pixel_mask.size())) {
        for (int i = 0; i < pixel_mask.size(); i++)
            pixel_mask[i] = start_msg.pixel_mask.begin()->second[i];
    }
}

void CBFWriter::WriteImage(const DataMessage &data_msg) {
    std::string filename = fmt::format("{:s}_{:06d}.cbf", start_msg.file_prefix, data_msg.number);
    FILE *fh = fopen(filename.c_str(), "wb");
    if (!fh)
        throw (JFJochException(JFJochExceptionCategory::FileWriteError, "Cannot create file " + filename));

    auto npixel = data_msg.image.GetWidth() * data_msg.image.GetHeight();

    std::vector<uint8_t> image;
    data_msg.image.GetUncompressed(image);

    std::vector<uint8_t> output(npixel * sizeof(uint32_t) * 2);

    switch (data_msg.image.GetMode()) {
        case CompressedImageMode::Int8:
            EncodeCBF((int8_t *) image.data(), npixel, output);
            break;
        case CompressedImageMode::Int16:
            EncodeCBF((int16_t *) image.data(), npixel, output);
            break;
        case CompressedImageMode::Int32:
            EncodeCBF((int32_t *) image.data(), npixel, output);
            break;
        case CompressedImageMode::Uint8:
            EncodeCBF((uint8_t *) image.data(), npixel, output);
            break;
        case CompressedImageMode::Uint16:
            EncodeCBF((uint16_t *) image.data(), npixel, output);
            break;
        case CompressedImageMode::Uint32:
            EncodeCBF((uint32_t *) image.data(), npixel, output);
            break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "RGB/float modes not supported in CBF writer");
    }

    fputs("###CBF: VERSION 1.7.11\r\n\r\n", fh);
    fputs(fmt::format("data_image_{:06d}\r\n\r\n", data_msg.number).c_str(), fh);
    fputs("_array_data.header_convention SLS_1.0\r\n", fh);
    fputs("_array_data.header_contents\r\n;\r\n\r\n", fh);
    fputs(header.c_str(), fh);

    fputs(";\r\n", fh);

    fputs("_array_data.data", fh);
    fputs("\r\n;\r\n", fh);
    fputs("--CIF-BINARY-FORMAT-SECTION--\r\n", fh);
    fputs("Content-Type: application/octet-stream;\r\n", fh);
    fputs("     conversions=\"x-CBF_BYTE_OFFSET\"\r\n", fh);
    fputs("Content-Transfer-Encoding: BINARY\r\n", fh);
    fputs(fmt::format("X-Binary-Size: {:d}\r\n", output.size()).c_str(), fh);
    fputs("X-Binary-ID: 1\r\n", fh);
    fputs("X-Binary-Element-Type: \"signed 32-bit integer\"\r\n", fh);
    fputs("X-Binary-Element-Byte-Order: LITTLE_ENDIAN\r\n", fh);
    // fputs("Content-MD5: {md5_hash}\r\n", fh);
    fputs(fmt::format("X-Binary-Number-of-Elements: {:d}\r\n", npixel).c_str(), fh);
    fputs(fmt::format("X-Binary-Size-Fastest-Dimension: {:d}\r\n", data_msg.image.GetWidth()).c_str(), fh);
    fputs(fmt::format("X-Binary-Size-Second-Dimension: {:d}\r\n", data_msg.image.GetHeight()).c_str(), fh);
    fputs("X-Binary-Size-Padding: 0\r\n", fh);

    fputs("\x0C\x1A\x04\xD5", fh);

    fwrite(output.data(), sizeof(uint8_t), output.size(), fh);
    fputs("\r\n--CIF-BINARY-FORMAT-SECTION----\r\n;\r\n\r\n",fh);

    fclose(fh);
}