Jungfraujoch/viewer/JFJochHttpReader.cpp

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

#include "JFJochHttpReader.h"

#include <cstring>
#include <mutex>
#include <nlohmann/json.hpp>
#include "../frame_serialize/CBORStream2Deserializer.h"
#include "../broker/gen/model/Image_buffer_status.h"
#include "../broker/gen/model/Plots.h"
#include "../broker/gen/model/Broker_status.h"
#include "../broker/gen/model/Roi_definitions.h"
#include "../common/InstrumentMetadata.h"
#include "../common/JFJochMath.h"
#include "../image_analysis/bragg_integration/CalcISigma.h"

// Included last, after all std/project headers: on Windows <httplib.h> pulls in <windows.h>,
// whose min/max (and other) macros must not precede the headers that use std::min / std::max.
#include <httplib.h>

JFJochHttpReader::HttpResult JFJochHttpReader::Request(const std::string &method, const std::string &path,
                                                       const std::string &body,
                                                       const std::string &content_type) const {
    // One client is reused across requests so the connection stays alive; if the socket has since
    // been closed, httplib reconnects on its own.
    std::lock_guard cl(client_mutex);

    if (!client) {
        // addr is scheme://host[:port][/prefix]; the client takes everything before the prefix.
        const auto host_start = addr.find("://");
        const auto prefix_start = addr.find('/', host_start == std::string::npos ? 0 : host_start + 3);
        path_prefix = (prefix_start == std::string::npos) ? "" : addr.substr(prefix_start);

        auto new_client = std::make_unique<httplib::Client>(addr.substr(0, prefix_start));
        if (!new_client->is_valid())
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Invalid broker address " + addr);
        client = std::move(new_client);
        client->set_keep_alive(true);
        // Cap connect and the whole request at 5 s so a stalled/black-holed broker fails the request
        // instead of blocking forever (which would wedge the reader thread and Close() on
        // client_mutex). All viewer requests are small and against a nearby broker, so 5 s is ample.
        client->set_connection_timeout(5);
        client->set_max_timeout(std::chrono::seconds(5));
    }

    httplib::Headers headers;
    if (!bearer_token.empty())
        headers.emplace("Authorization", "Bearer " + bearer_token);

    const std::string url = path_prefix + path;
    auto res = (method == "PUT")
                   ? client->Put(url, headers, body, content_type)
                   : client->Get(url, headers);

    HttpResult result;
    if (res) {
        result.ok = true;
        result.status = res->status;
        result.body = std::move(res->body);
    }

    // A protected dataset is reported the same way from every endpoint, so no caller has to know
    // which ones the broker guards.
    if (result.ok && (result.status == 401 || result.status == 403))
        throw HttpUnauthorized("Broker requires a token for the current dataset");
    return result;
}

void JFJochHttpReader::Token(const std::string &token) {
    std::unique_lock ul(http_mutex);
    bearer_token = token;
}

JFJochHttpReader::JFJochHttpReader() = default;

JFJochHttpReader::~JFJochHttpReader() {
    ResetConnection();
}

void JFJochHttpReader::ResetConnection() const {
    std::lock_guard cl(client_mutex);
    client.reset();
}

void JFJochHttpReader::Close() {
    std::unique_lock ul(http_mutex);
    addr = "";
    SetStartMessage({});
    last_image_buffer_counter = {};
    last_op_http_sync = false;
    cached_pixel_mask.reset();
    cached_pixel_mask_arm_date.clear();
    ResetConnection();  // drop the kept-alive connection on disconnect
}

ImageBufferStatus JFJochHttpReader::GetImageBufferStatus() const {
    auto res = Request("GET", "/image_buffer/status");
    if (!res.ok || res.status != 200)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Could not get image buffer status");

    try {
        org::openapitools::server::model::Image_buffer_status status = nlohmann::json::parse(res.body);
        ImageBufferStatus ret{};
        ret.max_image_number = status.getMaxImageNumber();
        ret.min_image_number = status.getMinImageNumber();
        ret.available_slots = status.getAvailableSlots();
        ret.total_slots = status.getTotalSlots();
        ret.images_in_the_buffer = status.getImageNumbers();
        if (status.currentCounterIsSet())
            ret.current_counter = status.getCurrentCounter();
        return ret;
    } catch (std::exception &e) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Could not parse image buffer status");
    }
}

BrokerStatus JFJochHttpReader::GetBrokerStatus() const {
    auto res = Request("GET", "/status");
    if (!res.ok || res.status != 200)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Could not get broker status");

    try {
        org::openapitools::server::model::Broker_status input = nlohmann::json::parse(res.body);
        BrokerStatus ret{};
        ret.broker_version = input.getBrokerVersion();
        ret.gpu_count = input.getGpuCount();
        if (input.progressIsSet())
            ret.progress = input.getProgress();
        if (input.messageIsSet())
            ret.message = input.getMessage();

        if (input.getState() == "Inactive")
            ret.state = JFJochState::Inactive;
        else if (input.getState() == "Idle")
            ret.state = JFJochState::Idle;
        else if (input.getState() == "Measuring")
            ret.state = JFJochState::Measuring;
        else if (input.getState() == "Error")
            ret.state = JFJochState::Error;
        else if (input.getState() == "Busy")
            ret.state = JFJochState::Busy;
        else if (input.getState() == "Pedestal")
            ret.state = JFJochState::Calibration;

        if (input.getMessageSeverity() == "info")
            ret.message_severity = BrokerStatus::MessageSeverity::Info;
        else if (input.getMessageSeverity() == "success")
            ret.message_severity = BrokerStatus::MessageSeverity::Success;
        else if (input.getMessageSeverity() == "warning")
            ret.message_severity = BrokerStatus::MessageSeverity::Warning;
        else if (input.getMessageSeverity() == "error")
            ret.message_severity = BrokerStatus::MessageSeverity::Error;
        if (input.brokerVersionIsSet())
            ret.broker_version = input.getBrokerVersion();

        return ret;
    } catch (std::exception &e) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Could not parse broker status");
    }
}

uint64_t JFJochHttpReader::GetNumberOfImages() const {
    std::unique_lock ul(http_mutex);

    if (addr.empty())
        return 0;
    auto status = GetImageBufferStatus();
    return status.max_image_number + 1;
}

std::shared_ptr<JFJochReaderDataset> JFJochHttpReader::UpdateDataset_i() {
    auto res = Request("GET", "/image_buffer/start.cbor");

    if (!res.ok || (res.status != 200 && res.status != 404))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Could not get image buffer status");

    if (res.status == 404)
        return {};
    if (res.body.empty())
        return {};

    try {
        auto msg = CBORStream2Deserialize(res.body);

        if (msg->msg_type != CBORImageType::START)
            return {};

        auto dataset = std::make_shared<JFJochReaderDataset>();
        dataset->arm_date = msg->start_message->arm_date;

        dataset->experiment = default_experiment;

        // JFJochReader is always using int32_t
        dataset->experiment.BitDepthImage(32);
        dataset->experiment.PixelSigned(true);
        dataset->experiment.FilePrefix(msg->start_message->file_prefix);
        dataset->experiment.BeamX_pxl(msg->start_message->beam_center_x);
        dataset->experiment.BeamY_pxl(msg->start_message->beam_center_y);
        dataset->experiment.DetectorDistance_mm(msg->start_message->detector_distance * 1000.0);
        dataset->experiment.DetectIceRings(msg->start_message->detect_ice_rings.value_or(false));
        dataset->experiment.PoniRot1_rad(msg->start_message->poni_rot1.value_or(0.0));
        dataset->experiment.PoniRot2_rad(msg->start_message->poni_rot2.value_or(0.0));
        dataset->experiment.PoniRot3_rad(msg->start_message->poni_rot3.value_or(0.0));

        dataset->az_int_bin_to_q = msg->start_message->az_int_bin_to_q;
        dataset->az_int_bin_to_phi = msg->start_message->az_int_bin_to_phi;
        dataset->q_bins = msg->start_message->az_int_q_bin_count.value_or(0);
        dataset->azimuthal_bins = msg->start_message->az_int_phi_bin_count.value_or(0);
        dataset->jfjoch_release = msg->start_message->jfjoch_release;

        DetectorSetup detector = DetDECTRIS(msg->start_message->image_size_x, msg->start_message->image_size_y,
                                            msg->start_message->detector_description, {});
        detector.PixelSize_um(msg->start_message->pixel_size_x * 1e6);
        detector.ImageOrientation(DetectorOrientation(msg->start_message->detector_orientation_mirror_y,
                                                      msg->start_message->detector_orientation_quarter_turns));
        detector.SaturationLimit(SaturationLimitFromValue(msg->start_message->saturation_value));
        detector.MinFrameTime(std::chrono::microseconds(0));
        detector.MinCountTime(std::chrono::microseconds(0));
        detector.ReadOutTime(std::chrono::microseconds (0));
        dataset->experiment.Detector(detector);
        // Absent (empty) for a flat detector, the image being the detector plane
        dataset->experiment.Modules(DetectorModules(msg->start_message->detector_modules));

        dataset->experiment.FrameTime(
                std::chrono::microseconds(std::lround(msg->start_message->frame_time * 1e6)),
                std::chrono::microseconds(std::lround(msg->start_message->count_time * 1e6))
        );

        if (!msg->start_message->pixel_mask.empty()) {
            // The pixel mask is constant for an acquisition; build the full-detector (~tens of MB)
            // PixelMask once per arm and share it across the per-refresh dataset snapshots, so a live
            // dataset refresh does not reconstruct and copy it on every tick.
            if (!cached_pixel_mask || cached_pixel_mask_arm_date != dataset->arm_date) {
                cached_pixel_mask =
                    std::make_shared<const PixelMask>(msg->start_message->pixel_mask.begin()->second);
                cached_pixel_mask_arm_date = dataset->arm_date;
            }
            dataset->pixel_mask = cached_pixel_mask;
        }

        dataset->experiment.NumTriggers(1);
        dataset->experiment.ImagesPerTrigger(msg->start_message->number_of_images);

        InstrumentMetadata instrument;
        instrument.SourceName(msg->start_message->source_name);
        instrument.SourceType(msg->start_message->source_type);
        instrument.InstrumentName(msg->start_message->instrument_name);
        dataset->experiment.ImportInstrumentMetadata(instrument);

        dataset->experiment.SampleName(msg->start_message->sample_name);
        dataset->experiment.SampleTemperature_K(msg->start_message->sample_temperature_K);
        dataset->experiment.RingCurrent_mA(msg->start_message->ring_current_mA);
        dataset->experiment.TotalFlux(msg->start_message->total_flux);
        dataset->experiment.AttenuatorTransmission(msg->start_message->attenuator_transmission);
        dataset->experiment.IncidentEnergy_keV(msg->start_message->incident_energy / 1000.0);

        dataset->experiment.FluorescenceSpectrum(msg->start_message->fluorescence_spectrum);

        dataset->bkg_estimate = GetPlot_i("bkg_estimate");
        dataset->spindle_blind_fraction = GetPlot_i("spindle_blind_fraction");
        dataset->ice_ring_score = GetPlot_i("ice_ring_score");
        dataset->spot_count = GetPlot_i("spot_count");
        dataset->spot_count_ice_rings = GetPlot_i("spot_count_ice");
        dataset->spot_count_low_res = GetPlot_i("spot_count_low_res");
        dataset->spot_count_indexed = GetPlot_i("spot_count_indexed");

        dataset->indexing_result = GetPlot_i("indexing_rate");
        dataset->indexing_lattice_count = GetPlot_i("indexing_lattice_count");
        dataset->profile_radius = GetPlot_i("profile_radius");
        dataset->mosaicity_deg = GetPlot_i("mosaicity");
        dataset->b_factor = GetPlot_i("b_factor");
        dataset->resolution_estimate = GetPlot_i("resolution_estimate");
        dataset->efficiency = GetPlot_i("image_collection_efficiency");
        dataset->integrated_reflections = GetPlot_i("integrated_reflections");
        dataset->image_scale_factor = GetPlot_i("image_scale_factor");
        dataset->image_scale_cc = GetPlot_i("image_scale_cc");

        // Not alternatives: a grid scan is collected at a head position the axis still reports, so
        // the START message carries both. Taking the axis and dropping the grid left the raster
        // drawn as a line plot (OpenAPIConvert.cpp had the same bug on the broker side).
        dataset->experiment.Goniometer(msg->start_message->goniometer);
        dataset->experiment.GridScan(msg->start_message->grid_scan);

        dataset->experiment.Smargon(msg->start_message->smargon_position);

        return dataset;
    } catch (std::exception &e) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              std::string("Could not load dataset: ") + std::string(e.what()));
    }
}

void JFJochHttpReader::ReadURL(const std::string &url) {
    std::unique_lock ul(http_mutex);
    ResetConnection();  // selecting an address (re)opens the connection from scratch
    // Drop the previous source's cached mask: a different (or re-armed) source has its own mask, and
    // the arm_date key alone does not distinguish an empty/duplicate arm_date across sources.
    cached_pixel_mask.reset();
    cached_pixel_mask_arm_date.clear();
    addr = url;
    if (url.empty())
        SetStartMessage({});
    else
        SetStartMessage(UpdateDataset_i());
}

std::shared_ptr<const JFJochReaderDataset> JFJochHttpReader::RefreshDatasetIfChanged(int64_t &num_images_out) {
    std::unique_lock ul(http_mutex);

    num_images_out = 0;
    if (addr.empty())
        return {};

    auto status = GetImageBufferStatus();
    num_images_out = status.max_image_number + 1;

    // Re-fetch the dataset (start message + plots) only when the buffer actually changed,
    // so the dataset-only follow mode does not hammer the broker with plot requests.
    const bool buffer_changed = !last_image_buffer_counter.has_value()
        || !status.current_counter.has_value()
        || last_image_buffer_counter.value() != status.current_counter.value();
    last_image_buffer_counter = status.current_counter;

    if (!buffer_changed)
        return {};

    SetStartMessage(UpdateDataset_i());
    return GetDataset();
}

bool JFJochHttpReader::LoadImage_i(std::shared_ptr<JFJochReaderDataset> &dataset,
                                   DataMessage &message,
                                   std::vector<uint8_t> &buffer,
                                   int64_t image_number,
                                   bool update_dataset) {
    std::unique_lock ul(http_mutex);

    if (addr.empty())
        return false;

    bool buffer_changed = false;

    // For autoupdate - if buffer didn't change don't update dataset
    auto status = GetImageBufferStatus();
    if (!last_image_buffer_counter.has_value() // No information on the previous buffer state - always assume it changed
        || !status.current_counter.has_value() // No information on the current buffer state - e.g. old version of software
        || last_image_buffer_counter.value() != status.current_counter.value()) // current counter value different from previous
        buffer_changed = true;
    last_image_buffer_counter = status.current_counter;

    if (image_number == -1) {
        if (last_op_http_sync && !buffer_changed)
            return false;
        last_op_http_sync = true;
    } else {
        last_op_http_sync = false;
    }

    // Always update dataset, as it might have changed from the last time
    if (buffer_changed && update_dataset)
        dataset = UpdateDataset_i();
    if (!dataset)
        return false;

    auto res = Request("GET", "/image_buffer/image.cbor?id=" + std::to_string(image_number));

    if (!res.ok || res.status != 200)
        return false;

    if (res.body.empty()) {
        return false;
    }

    try {
        buffer.resize(res.body.size());
        memcpy(buffer.data(), res.body.data(), res.body.size());

        auto msg = CBORStream2Deserialize(buffer);

        if (msg->msg_type != CBORImageType::IMAGE)
            return false;

        message = *msg->data_message;

        CalcISigma(message);
        CalcWilsonBFactor(message);

        return true;
    } catch (std::exception &e) {
        return false;
    }
}

bool JFJochHttpReader::ReadRawImage(int64_t image_number, JFJochReaderRawImage &image) {
    if (addr.empty())
        return false;

    HttpResult res;
    try {
        res = Request("GET", "/image_buffer/image.cbor?id=" + std::to_string(image_number));
    } catch (const HttpUnauthorized &) {
        return false;
    }

    if (!res.ok || res.status != 200 || res.body.empty())
        return false;

    try {
        auto msg =
                CBORStream2Deserialize(reinterpret_cast<uint8_t *>(res.body.data()), res.body.size());

        if (msg->msg_type != CBORImageType::IMAGE)
            return false;

        image.image_buffer.resize(msg->data_message->image.GetCompressedSize());
        memcpy(image.image_buffer.data(),
               msg->data_message->image.GetCompressed(),
               msg->data_message->image.GetCompressedSize());
        image.image = CompressedImage(
            image.image_buffer.data(),
            image.image_buffer.size(),
            msg->data_message->image.GetWidth(),
            msg->data_message->image.GetHeight(),
            msg->data_message->image.GetMode(),
            msg->data_message->image.GetCompressionAlgorithm()
        );

        return true;
    } catch (std::exception &e) {
        return false;
    }
}

std::vector<float> JFJochHttpReader::GetPlot_i(const std::string &plot_type, float fill_value) const {
    if (addr.empty())
        return {};

    auto res_bin = Request("GET", "/preview/plot.bin?type=" + plot_type);
    if (res_bin.ok && res_bin.status == 200) {
        if (res_bin.body.size() % sizeof(float) != 0) {
            throw std::runtime_error("Input size is not a multiple of sizeof(float)");
        }
        std::vector<float> v(res_bin.body.size() / sizeof(float));
        std::memcpy(v.data(), res_bin.body.data(), res_bin.body.size());
        return v;
    }

    auto res = Request("GET", "/preview/plot?binning=1&experimental_coord=false&fill="
                       + std::to_string(fill_value) + "&type=" + plot_type);

    if (!res.ok || res.status != 200)
        return {};

    try {
        org::openapitools::server::model::Plots plots = nlohmann::json::parse(res.body);
        auto plot_v = plots.getPlot();
        if (plot_v.size() == 1)
            return plot_v[0].getY();
        else
            return {};
    } catch (nlohmann::json::parse_error &e) {
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Could not parse plot " + plot_type);
    }
}

void JFJochHttpReader::UploadUserMask(const std::vector<uint32_t>& mask) {
    std::unique_lock ul(http_mutex);

    if (addr.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "HTTP address not set. Call ReadURL() first.");

    if (mask.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "User mask is empty.");

    const char* data_ptr = reinterpret_cast<const char*>(mask.data());
    const size_t byte_size = mask.size() * sizeof(uint32_t);

    auto res = Request("PUT", "/config/user_mask",
                       std::string(data_ptr, byte_size),
                       "application/octet-stream");

    if (!res.ok)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Failed to connect to server to upload user mask");

    if (res.status != 200)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Server rejected user mask upload");
}

ROIDefinition JFJochHttpReader::GetROIDefinitions() const {
    std::unique_lock ul(http_mutex);
    if (addr.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "HTTP address not set");

    auto res = Request("GET", "/config/roi");
    if (!res.ok || res.status != 200)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Could not get ROI definitions");

    org::openapitools::server::model::Roi_definitions input = nlohmann::json::parse(res.body);
    ROIDefinition out;
    for (const auto &i : input.getBox().getRois())
        out.boxes.emplace_back(i.getName(), i.getMinXPxl(), i.getMaxXPxl(), i.getMinYPxl(), i.getMaxYPxl());
    for (const auto &i : input.getCircle().getRois())
        out.circles.emplace_back(i.getName(), i.getCenterXPxl(), i.getCenterYPxl(), i.getRadiusPxl());
    for (const auto &i : input.getAzim().getRois()) {
        float phi_min = 0, phi_max = 0;
        if (i.phiMinDegIsSet() && i.phiMaxDegIsSet()) {
            phi_min = i.getPhiMinDeg();
            phi_max = i.getPhiMaxDeg();
        }
        const float d_min = (i.getQMaxRecipA() == 0.0f) ? 0.0f : 2.0f * static_cast<float>(PI) / i.getQMaxRecipA();
        const float d_max = (i.getQMinRecipA() == 0.0f) ? 0.0f : 2.0f * static_cast<float>(PI) / i.getQMinRecipA();
        out.azimuthal.emplace_back(i.getName(), d_min, d_max, phi_min, phi_max);
    }
    return out;
}

void JFJochHttpReader::UploadROIDefinitions(const ROIDefinition &rois) const {
    std::unique_lock ul(http_mutex);
    if (addr.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "HTTP address not set");

    namespace model = org::openapitools::server::model;
    model::Roi_definitions out;

    model::Roi_box_list bl;
    std::vector<model::Roi_box> boxes;
    for (const auto &b : rois.boxes) {
        model::Roi_box e;
        e.setName(b.GetName());
        e.setMinXPxl(b.GetXMin()); e.setMaxXPxl(b.GetXMax());
        e.setMinYPxl(b.GetYMin()); e.setMaxYPxl(b.GetYMax());
        boxes.push_back(e);
    }
    bl.setRois(boxes);
    out.setBox(bl);

    model::Roi_circle_list cl;
    std::vector<model::Roi_circle> circles;
    for (const auto &c : rois.circles) {
        model::Roi_circle e;
        e.setName(c.GetName());
        e.setCenterXPxl(c.GetX()); e.setCenterYPxl(c.GetY());
        e.setRadiusPxl(c.GetRadius_pxl());
        circles.push_back(e);
    }
    cl.setRois(circles);
    out.setCircle(cl);

    model::Roi_azim_list al;
    std::vector<model::Roi_azimuthal> azim;
    for (const auto &a : rois.azimuthal) {
        model::Roi_azimuthal e;
        e.setName(a.GetName());
        e.setQMinRecipA(a.GetQMin_recipA());
        e.setQMaxRecipA(a.GetQMax_recipA());
        if (a.HasPhi()) {
            e.setPhiMinDeg(a.GetPhiMin_deg());
            e.setPhiMaxDeg(a.GetPhiMax_deg());
        }
        azim.push_back(e);
    }
    al.setRois(azim);
    out.setAzim(al);

    nlohmann::json j = out;
    auto res = Request("PUT", "/config/roi", j.dump(), "application/json");
    if (!res.ok)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Failed to connect to server to upload ROIs");
    if (res.status != 200)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Server rejected ROI upload");
}

std::vector<SpotToSave> JFJochHttpReader::ReadSpots(int64_t image) const {
    std::unique_lock ul(http_mutex);

    if (image < 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
            "Image number must be non-negative");

    if (addr.empty())
        return {};

    HttpResult res;
    try {
        res = Request("GET", "/image_buffer/image.cbor?id=" + std::to_string(image));
    } catch (const HttpUnauthorized &) {
        return {};
    }

    if (!res.ok || res.status != 200 || res.body.empty())
        return {};

    try {
        auto msg = CBORStream2Deserialize(res.body);

        if (msg->msg_type != CBORImageType::IMAGE)
            return {};

        return msg->data_message->spots;
    } catch (std::exception &e) {
        return {};
    }
}