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

#ifdef JFJOCH_USE_NUMA
#include <numaif.h>
#endif

#include <sys/mman.h>
#include <thread>
#include <fstream>
#include <cmath>

#include "../common/JFJochException.h"
#include "AcquisitionDevice.h"
#include "../common/NetworkAddressConvert.h"

void *mmap_acquisition_buffer(size_t size, int16_t numa_node) {
    void *ret = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (ret == MAP_FAILED) {
        throw JFJochException(JFJochExceptionCategory::MemAllocFailed, "frame_buffer");
    }
#ifdef JFJOCH_USE_NUMA
    if (numa_node >= 0) {
        unsigned long nodemask = 1L << numa_node;;
        if (numa_node > sizeof(nodemask)*8)
            throw JFJochException(JFJochExceptionCategory::MemAllocFailed, "Mask too small for NUMA node");
        if (mbind(ret, size, MPOL_BIND, &nodemask, sizeof(nodemask)*8, MPOL_MF_STRICT) == -1)
            throw JFJochException(JFJochExceptionCategory::MemAllocFailed, "Cannot apply NUMA policy");
    }
#endif
    memset(ret, 0, size);
    return ret;
}

AcquisitionDevice::AcquisitionDevice(uint16_t in_data_stream) {
    logger = nullptr;
    data_stream = in_data_stream;
}

void AcquisitionDevice::PrepareAction(const DiffractionExperiment &experiment) {
    if (experiment.GetModulesNum(data_stream) > max_modules)
        throw(JFJochException(JFJochExceptionCategory::InputParameterAboveMax,
                              "Number of modules exceeds max possible for FPGA"));

    counters.Reset(experiment, data_stream);
}

void AcquisitionDevice::StartAction(const DiffractionExperiment &experiment, uint32_t optional_flags) {
    Cancel();

    if (experiment.GetModulesNum(data_stream) > max_modules)
        throw(JFJochException(JFJochExceptionCategory::InputParameterAboveMax,
                              "Number of modules exceeds max possible for FPGA"));

    counters.Reset(experiment, data_stream);
    expected_frames = experiment.GetFrameNum() / experiment.GetFPGASummation();

    // Ensure internal WR queue is empty
    work_request_queue.Clear();

    Start(experiment, optional_flags);

    for (uint32_t i = 0; i < buffer_device.size(); i++)
        SendWorkRequest(i);

    auto c = work_completion_queue.GetBlocking();
    if (c.type != Completion::Type::Start)
        throw JFJochException(JFJochExceptionCategory::AcquisitionDeviceError, "Mismatch in work completions");

    StartSendingWorkRequests();

    start_time = std::chrono::system_clock::now();

    if (experiment.IsUsingInternalPacketGen())
        RunInternalGenerator(experiment);
}

void AcquisitionDevice::WaitForActionComplete() {
    auto c = work_completion_queue.GetBlocking();

    while (c.type != Completion::Type::End) {
        DeviceOutput* output;

        try {
            output = GetDeviceOutput(c.handle);
        } catch (const JFJochException &e) {
            if (logger)
                logger->ErrorException(e);
            continue;
        }

        c.module_number = output->module_statistics.module_number;
        c.packet_count = output->module_statistics.packet_count;
        c.frame_number = output->module_statistics.frame_number;

        if (c.frame_number >= expected_frames) {
            Cancel();
            // this frame is not of any interest, therefore its location can be immediately released
            SendWorkRequest(c.handle);
        } else if (c.module_number >= max_modules) {
            // Module number out of bounds, don't process
            if (logger != nullptr)
                logger->Error("Completion with wrong module number data stream {} completion frame number {} module {} handle {}",
                              data_stream, c.frame_number, c.module_number, c.handle);
            SendWorkRequest(c.handle);
        } else if (c.frame_number < counters.GetSlowestFrameNumber()) {
            // Module is falling behind, needs to return the handle then
            SendWorkRequest(c.handle);
        } else {
            try {
                counters.UpdateCounters(&c);
            } catch (const JFJochException &e) {
                if (logger)
                    logger->ErrorException(e);
                SendWorkRequest(c.handle);
            }
        }
        if (logger != nullptr)
            logger->Debug("Data stream {} completion frame number {} module {} handle {}",
                          data_stream, c.frame_number, c.module_number, c.handle);

        c = work_completion_queue.GetBlocking();
    }
    counters.SetAcquisitionFinished();

    end_time = std::chrono::system_clock::now();
    Cancel();
    Finalize();
}

void AcquisitionDevice::SendWorkRequest(uint32_t handle) {
    work_request_queue.Put(WorkRequest{
        .handle = handle
    });
}

uint64_t AcquisitionDevice::GetBytesReceived() const {
    return counters.GetBytesReceived();
}

const DeviceOutput *AcquisitionDevice::GetDeviceOutput(size_t frame_number, uint16_t module_number) const {
    auto handle = counters.GetBufferHandle(frame_number, module_number);
    if (handle != HandleNotValid)
        return GetDeviceOutput(handle);
    else
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Frame not collected");
}

const DeviceOutput *AcquisitionDevice::GetDeviceOutput(size_t handle) const {
    if (handle >= buffer_device.size())
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Handle outside of range");
    else
        return (DeviceOutput *) buffer_device.at(handle);
}

DeviceOutput *AcquisitionDevice::GetDeviceOutput(size_t handle) {
    if (handle >= buffer_device.size())
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds, "Handle outside of range");
    else
        return (DeviceOutput *) buffer_device.at(handle);
}

void AcquisitionDevice::InitializeCalibration(const DiffractionExperiment &experiment, const JFCalibration &calib) {}

void AcquisitionDevice::InitializeIntegrationMap(const DiffractionExperiment &experiment,
                                                 const std::vector<uint16_t> &v,
                                                 const std::vector<float> &weights) {}

void AcquisitionDevice::InitializeIntegrationMap(const uint16_t *map, const float *weights, size_t module_number) {}

void AcquisitionDevice::InitializeSpotFinderResolutionMap(const float *data, size_t module_number) {}

void AcquisitionDevice::InitializeROIMap(const uint16_t *map, size_t module_number) {}

void AcquisitionDevice::InitializePixelMask(const uint32_t *module_mask, size_t module_number) {}

void AcquisitionDevice::InitializeROIMap(const DiffractionExperiment& experiment, const std::vector<uint16_t>& roi_map) {
    if (roi_map.size() != experiment.GetXPixelsNumConv() * experiment.GetYPixelsNumConv())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in array size");

    std::vector<uint16_t> tmp(RAW_MODULE_SIZE);
    auto offset = experiment.GetFirstModuleOfDataStream(data_stream);
    size_t modules = experiment.GetModulesNum(data_stream);
    for (int m = 0; m < modules; m++) {
        ConvertedToRawGeometry(experiment, offset + m, tmp.data(), roi_map.data());
        InitializeROIMap(tmp.data(), m);
    }
}

void AcquisitionDevice::InitializeEmptyPixelMask(const DiffractionExperiment &experiment) {
    std::vector<uint32_t> empty_mask(RAW_MODULE_SIZE);

    size_t modules = experiment.GetModulesNum(data_stream);
    for (int m = 0; m < modules; m++)
        InitializePixelMask(empty_mask.data(), m);
}


void AcquisitionDevice::InitializeDataProcessing(const DiffractionExperiment &experiment,
                                                 const AzimuthalIntegration &azint) {
    auto offset = experiment.GetFirstModuleOfDataStream(data_stream);
    size_t modules = experiment.GetModulesNum(data_stream);

    if (experiment.IsGeometryTransformed()) {
        std::vector<float> tmp1(RAW_MODULE_SIZE);
        std::vector<uint16_t> tmp2(RAW_MODULE_SIZE);

        for (int m = 0; m < modules; m++) {
            ConvertedToRawGeometry(experiment, offset + m, tmp1.data(), azint.Corrections().data());
            ConvertedToRawGeometry(experiment, offset + m, tmp2.data(), azint.GetPixelToBin().data());
            InitializeIntegrationMap(tmp2.data(), tmp1.data(), m);

            ConvertedToRawGeometry(experiment, offset + m, tmp1.data(), azint.Resolution().data());
            InitializeSpotFinderResolutionMap(tmp1.data(), m);
        }
    } else {
        for (int m = 0; m < modules; m++) {
            InitializeIntegrationMap(azint.GetPixelToBin().data() + (offset + m) * RAW_MODULE_SIZE,
                                     azint.Corrections().data() + (offset + m) * RAW_MODULE_SIZE,
                                     m);
            InitializeSpotFinderResolutionMap(azint.Resolution().data() + (m + offset) * RAW_MODULE_SIZE,
                                              m);
        }
    }
}

void AcquisitionDevice::InitializePixelMask(const DiffractionExperiment &experiment, const PixelMask &mask) {
    auto offset = experiment.GetFirstModuleOfDataStream(data_stream);
    size_t modules = experiment.GetModulesNum(data_stream);

    std::vector<uint32_t> tmp(RAW_MODULE_SIZE);
    for (int m = 0; m < modules; m++) {
        ConvertedToRawGeometry(experiment, offset + m, tmp.data(), mask.GetMask().data());
        InitializePixelMask(tmp.data(), m);
    }
}

void AcquisitionDevice::MapBuffersStandard(size_t c2h_buffer_count, int16_t numa_node) {
    try {
        for (int i = 0; i < c2h_buffer_count; i++)
            buffer_device.emplace_back((DeviceOutput *) mmap_acquisition_buffer(FPGA_BUFFER_LOCATION_SIZE, numa_node));
    } catch (const JFJochException &e) {
        UnmapBuffers();
        throw;
    }
}

void AcquisitionDevice::UnmapBuffers() {
     for (auto &i: buffer_device)
        if (i != nullptr) munmap(i, FPGA_BUFFER_LOCATION_SIZE);
}

void AcquisitionDevice::FrameBufferRelease(size_t frame_number, uint16_t module_number) {
    auto handle = counters.GetBufferHandleAndClear(frame_number, module_number);
    if (handle != AcquisitionCounters::HandleNotFound)
        SendWorkRequest(handle);
}

void AcquisitionDevice::EnableLogging(Logger *in_logger) {
    logger = in_logger;
}

int32_t AcquisitionDevice::GetNUMANode() const {
    return -1;
}

uint16_t AcquisitionDevice::GetUDPPort() const {
    return 1234;
}

const AcquisitionCounters &AcquisitionDevice::Counters() const {
    return counters;
}

std::string AcquisitionDevice::GetIPv4Address() const {
    return IPv4AddressToStr(ipv4_addr);
}

std::string AcquisitionDevice::GetMACAddress() const {
    return MacAddressToStr(mac_addr);
}

DataCollectionStatus AcquisitionDevice::GetDataCollectionStatus() const {
    return {};
}

DeviceStatus AcquisitionDevice::GetDeviceStatus() const {
    return {};
}

AcquisitionDeviceStatistics AcquisitionDevice::GetStatistics() const {
    AcquisitionDeviceStatistics ret{};
    ret.bytes_received = GetBytesReceived();
    ret.start_timestamp = start_time.time_since_epoch().count();
    ret.end_timestamp = end_time.time_since_epoch().count();
    ret.packets_expected = counters.GetTotalExpectedPackets();
    ret.good_packets = counters.GetTotalPackets();

    for (int i = 0; i < counters.GetModuleNumber(); i++)
        ret.packets_received_per_module.push_back(counters.GetTotalPackets(i));

    if ((ret.packets_expected == 0) || (ret.good_packets == ret.packets_expected))
        ret.efficiency = 1.0;
    else
        ret.efficiency = static_cast<float>(ret.good_packets) / static_cast<float>(ret.packets_expected);

    return ret;
}

void AcquisitionDevice::SetIPv4Address(uint32_t ipv4_addr_network_order) {
    ipv4_addr = ipv4_addr_network_order;
}

AcquisitionDeviceNetConfig AcquisitionDevice::GetNetConfig() const {
    return {
        .mac_addr = GetMACAddress(),
        .ipv4_addr = GetIPv4Address(),
        .udp_port = GetUDPPort()
    };
}

void AcquisitionDevice::RunInternalGenerator(const DiffractionExperiment &experiment) {
    FrameGeneratorConfig config{};
    config.frames = experiment.GetFrameNum() + DELAY_FRAMES_STOP_AND_QUIT + 1;
    config.modules = experiment.GetModulesNum(data_stream);
    config.data_stream = data_stream;
    config.pulse_id = INT_PKT_GEN_BUNCHID;
    config.exptime = INT_PKT_GEN_EXPTTIME;
    config.debug = INT_PKT_GEN_DEBUG;
    config.dest_mac_addr = MacAddressFromStr(GetMACAddress());
    config.dest_ipv4_addr = IPv4AddressFromStr(GetIPv4Address());
    config.images_in_memory = experiment.GetInternalPacketGeneratorImages() - 1;
    switch (experiment.GetDetectorSetup().GetDetectorType()) {
        case DetectorType::JUNGFRAU:
            config.detector_type = SLS_DETECTOR_TYPE_JUNGFRAU;
            break;
        case DetectorType::EIGER:
            config.detector_type = SLS_DETECTOR_TYPE_EIGER;
            config.eiger_bit_depth = experiment.GetBitDepthReadout();
            break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Detector not supported");
    }
    HW_RunInternalGenerator(config);
}

void AcquisitionDevice::SetSpotFinderParameters(const SpotFindingSettings &settings) {
    SpotFinderParameters fpga_parameters{};

    fpga_parameters.snr_threshold = settings.signal_to_noise_threshold;
    fpga_parameters.count_threshold = settings.photon_count_threshold;
    fpga_parameters.max_d = settings.low_resolution_limit;
    fpga_parameters.min_d = settings.high_resolution_limit;
    fpga_parameters.min_pix_per_spot = settings.min_pix_per_spot;
    HW_SetSpotFinderParameters(fpga_parameters);
}