// Copyright (2019-2023) Paul Scherrer Institute

#include "JFJochReceiverService.h"

PlotRequest Convert(const JFJochProtoBuf::PlotRequest& request) {
    PlotRequest ret;
    ret.binning = request.binning();
    switch (request.type()) {
        case JFJochProtoBuf::BKG_ESTIMATE:
            ret.type = PlotType::BkgEstimate;
            break;
        case JFJochProtoBuf::RAD_INT:
            ret.type = PlotType::RadInt;
            break;
        case JFJochProtoBuf::SPOT_COUNT:
            ret.type = PlotType::SpotCount;
            break;
        case JFJochProtoBuf::INDEXING_RATE:
            ret.type = PlotType::IndexingRate;
            break;
        case JFJochProtoBuf::INDEXING_RATE_PER_FILE:
            ret.type = PlotType::IndexingRatePerFile;
            break;
        default:
        case JFJochProtoBuf::ADU_HISTOGRAM:
            ret.type = PlotType::ADUHistorgram;
            break;
    }
    return ret;
}

void Convert(const Plot& input, JFJochProtoBuf::Plot &output) {
    if (!input.x.empty())
        *output.mutable_x() = {input.x.begin(), input.x.end()};
    if (!input.y.empty())
        *output.mutable_y() = {input.y.begin(), input.y.end()};

}

void Convert(const RadialIntegrationProfiles& input, JFJochProtoBuf::RadialIntegrationProfiles& output) {
    for (const auto &i: input.profiles) {
        auto tmp = output.add_profiles();
        tmp->set_title(i.title);
        Convert(i.plot, *tmp->mutable_plot());
    }
}

JFJochReceiverService::JFJochReceiverService(std::vector<AcquisitionDevice *> &open_capi_device,
                                             Logger &in_logger, ImagePusher &pusher) :
        logger(in_logger), aq_devices(open_capi_device),
        image_pusher(pusher), data_processing_settings(DiffractionExperiment::DefaultDataProcessingSettings()) {
}

grpc::Status JFJochReceiverService::Start(grpc::ServerContext *context, const JFJochProtoBuf::ReceiverInput *request,
                                          JFJochProtoBuf::Empty *response) {
    std::unique_lock<std::mutex> ul_state(state_mutex);
    if (state != ReceiverState::Idle)
        return {grpc::StatusCode::ABORTED, "Receiver not idle, cannot start"};

    try {
        // Thanks to properties of unique_ptr, starting new measurement will call destructor of JFJochReceiver, which will
        // ensure that everything was rolled back
        // But it is important to do it in correct order - first abort old receiver, close it, than start new one
        receiver.reset();
        experiment = std::make_unique<DiffractionExperiment>(request->jungfraujoch_settings());
        if (request->has_calibration())
            calibration = std::make_unique<JFCalibration>(request->calibration());
        else
            calibration.reset();

        receiver = std::make_unique<JFJochReceiver>(*experiment, calibration.get(),
                                                    aq_devices, image_pusher,
                                                    logger, nthreads, send_buffer_count,
                                                    preview_publisher, numa_policy);
        try {
            // Don't want to stop
            receiver->SetDataProcessingSettings(data_processing_settings);
        } catch (...) {}
        measurement = std::async(std::launch::async, &JFJochReceiverService::FinalizeMeasurement, this);
        state = ReceiverState::Running;

        return grpc::Status::OK;
    } catch (const JFJochException &e) {
        logger.ErrorException(e);
        return {grpc::StatusCode::ABORTED, e.what()};
    }

}

void JFJochReceiverService::FinalizeMeasurement() {
    receiver->StopReceiver();
    {
        std::unique_lock<std::mutex> ul(state_mutex);
        state = ReceiverState::Idle;
        measurement_done.notify_all();
    }
}

grpc::Status JFJochReceiverService::Stop(grpc::ServerContext *context, const JFJochProtoBuf::Empty *request,
                                         JFJochProtoBuf::ReceiverOutput *response) {
    std::unique_lock<std::mutex> ul(state_mutex);

    measurement_done.wait(ul, [this] { return (state != ReceiverState::Running);});

    if (state == ReceiverState::Idle) {
        try {
            if (measurement.valid())
                measurement.get();
        } catch (JFJochException &e) {
            logger.ErrorException(e);
            return {grpc::StatusCode::ABORTED, e.what()};
        }

        if (!receiver) {
            logger.Warning("Request via gRPC, while receiver not running");
            return {grpc::StatusCode::ABORTED, "Receiver not running"};
        } else {
            auto output = receiver->GetStatistics();
            response->set_max_receive_delay(output.max_receive_delay);
            response->set_compressed_size(output.compressed_size);
            response->set_compressed_ratio(output.compressed_ratio);
            response->set_images_sent(output.images_sent);
            response->set_start_time_ms(output.start_time_ms);
            response->set_end_time_ms(output.end_time_ms);
            response->set_efficiency(output.efficiency);
            response->set_max_image_number_sent(output.max_image_number_sent);
            response->set_cancelled(output.cancelled);
            response->set_master_file_name(output.master_file_name);
            if (!output.pedestal_result.empty())
                *response->mutable_pedestal_result() = {output.pedestal_result.begin(), output.pedestal_result.end()};
            response->set_indexing_rate(output.indexing_rate);
            response->set_bkg_estimate(output.bkg_estimate);
            for (const auto &i: output.received_packets)
                response->add_received_packets(i);
            for (const auto &i: output.expected_packets)
                response->add_expected_packets(i);
            return grpc::Status::OK;
        }
    } else
        return {grpc::StatusCode::ABORTED, "Strange state"};
}

grpc::Status JFJochReceiverService::Cancel(grpc::ServerContext *context, const JFJochProtoBuf::Empty *request,
                                          JFJochProtoBuf::Empty *response) {
    std::unique_lock<std::mutex> ul(state_mutex);
    if (state == ReceiverState::Idle)
        return {grpc::StatusCode::ABORTED, "Receiver not running"};
    else {
        receiver->Cancel();
        return grpc::Status::OK;
    }
}

grpc::Status JFJochReceiverService::Abort(grpc::ServerContext *context, const JFJochProtoBuf::Empty *request,
                                          JFJochProtoBuf::Empty *response) {
    std::unique_lock<std::mutex> ul(state_mutex);
    if (state == ReceiverState::Idle)
        return {grpc::StatusCode::ABORTED, "Receiver not running"};
    else {
        receiver->Cancel();
        return grpc::Status::OK;
    }
}

JFJochReceiverService& JFJochReceiverService::NumThreads(int64_t input) {
    if (input <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Thread number must be above zero");
    nthreads = input;
    return *this;
}

JFJochReceiverService &JFJochReceiverService::SendBufferCount(int64_t input) {
    if (input <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Send buffer count must be above zero");
    send_buffer_count = input;
    return *this;
}

JFJochReceiverService& JFJochReceiverService::PreviewPublisher(ZMQPreviewPublisher *in_preview_writer) {
    preview_publisher = in_preview_writer;
    return *this;
}

JFJochReceiverService &JFJochReceiverService::NUMAPolicy(const NUMAHWPolicy &policy) {
    numa_policy = policy;
    return *this;
}

JFJochReceiverService &JFJochReceiverService::NUMAPolicy(const std::string &policy) {
    numa_policy = NUMAHWPolicy(policy);
    return *this;
}

grpc::Status JFJochReceiverService::GetStatus(grpc::ServerContext *context, const JFJochProtoBuf::Empty *request,
                                                JFJochProtoBuf::ReceiverStatus *response) {
    // Need to hold mutex, as receiver might not exist here, if state is idle
    std::unique_lock<std::mutex> ul(state_mutex);

    response->set_progress(-1);
    response->set_send_buffers_avail(-1);
    response->set_indexing_rate(-1);
    if (state == ReceiverState::Running) {
        response->set_progress(receiver->GetProgress());
        response->set_send_buffers_avail(receiver->GetAvailableSendBuffers());
    }
    if (receiver) {
        double indexing_rate = receiver->GetIndexingRate();
        if (!std::isnan(indexing_rate))
            response->set_indexing_rate(indexing_rate);
    }

    return grpc::Status::OK;
}

grpc::Status JFJochReceiverService::GetDataProcessingPlots(grpc::ServerContext *context,
                                                           const JFJochProtoBuf::PlotRequest *request,
                                                           JFJochProtoBuf::Plot *response) {
    // Need to hold mutex, as receiver might not exist here, if state is idle
    std::unique_lock<std::mutex> ul(state_mutex);
    if (receiver)
        Convert(receiver->GetPlots(Convert(*request)), *response);
    return grpc::Status::OK;
}

grpc::Status JFJochReceiverService::GetRadialIntegrationProfiles(grpc::ServerContext *context,
                                                                 const JFJochProtoBuf::Empty *request,
                                                                 JFJochProtoBuf::RadialIntegrationProfiles *response) {
    // Need to hold mutex, as receiver might not exist here, if state is idle
    std::unique_lock<std::mutex> ul(state_mutex);
    if (receiver)
        Convert(receiver->GetRadialIntegrationProfiles(), *response);
    return grpc::Status::OK;
}

grpc::Status JFJochReceiverService::SetDataProcessingSettings(grpc::ServerContext *context,
                                                              const JFJochProtoBuf::DataProcessingSettings *request,
                                                              JFJochProtoBuf::Empty *response) {
    try {
        std::unique_lock<std::mutex> ul(state_mutex);
        data_processing_settings = *request;
        if (state != ReceiverState::Idle)
            receiver->SetDataProcessingSettings(*request);

        return grpc::Status::OK;
    } catch (std::exception &e) {
        return {grpc::StatusCode::ABORTED, e.what()};
    }
}

grpc::Status JFJochReceiverService::GetNetworkConfig(grpc::ServerContext *context, const JFJochProtoBuf::Empty *request,
                                                     JFJochProtoBuf::ReceiverNetworkConfig *response) {
    for (const auto &aq: aq_devices) {
        auto dev_net_cfg = response->add_device();
        dev_net_cfg->set_mac_addr(aq->GetMACAddress());
        dev_net_cfg->set_ipv4_addr(aq->GetIPv4Address());
        dev_net_cfg->set_udp_port(aq->GetUDPPort());
    }
    return grpc::Status::OK;
}